diff --git a/lib/std/treap.zig b/lib/std/treap.zig index 1c4541f96f..98e85928d5 100644 --- a/lib/std/treap.zig +++ b/lib/std/treap.zig @@ -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| {