zig

commit da7ecfb2dec3869d3367fc730b677ae404ecac60 (tree)
parent 2adb932ad6ee4ff3d3c640cb8fb7bf7db0ff5d74
Author: Ryan Liptak <squeek502@hotmail.com>
Date:   Mon,  2 Oct 2023 20:53:53 -0700

Treap: Add InorderIterator

Diffstat:
M
lib/std/treap.zig
 | 
52
++++++++++++++++++++++++++++++++++++++++++++++++++++

1 file changed, 52 insertions(+), 0 deletions(-)
diff --git 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| {