Treap: Add InorderIterator

This commit is contained in:
Ryan Liptak 2023-10-02 20:53:53 -07:00
parent 2adb932ad6
commit da7ecfb2de

View File

@ -257,6 +257,48 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {
assert(link.* == node);
link.* = target;
}
pub const InorderIterator = struct {
current: ?*Node,
previous: ?*Node = null,
pub fn next(it: *InorderIterator) ?*Node {
while (true) {
if (it.current) |current| {
const previous = it.previous;
it.previous = current;
if (previous == current.parent) {
if (current.children[0]) |left_child| {
it.current = left_child;
} else {
if (current.children[1]) |right_child| {
it.current = right_child;
} else {
it.current = current.parent;
}
return current;
}
} else if (previous == current.children[0]) {
if (current.children[1]) |right_child| {
it.current = right_child;
} else {
it.current = current.parent;
}
return current;
} else {
std.debug.assert(previous == current.children[1]);
it.current = current.parent;
}
} else {
return null;
}
}
}
};
pub fn inorderIterator(self: *Self) InorderIterator {
return .{ .current = self.root };
}
};
}
@ -344,6 +386,16 @@ test "std.Treap: insert, find, replace, remove" {
try testing.expectEqual(entry.node, treap.getEntryForExisting(node).node);
}
// in-order iterator check
{
var it = treap.inorderIterator();
var last_key: u64 = 0;
while (it.next()) |node| {
try std.testing.expect(node.key >= last_key);
last_key = node.key;
}
}
// replace check
iter.reset();
while (iter.next()) |node| {