| ... | @@ -0,0 +1,294 @@ |
| 1 | const std = @import("std.zig"); |
| 2 | const assert = std.debug.assert; |
| 3 | const testing = std.testing; |
| 4 | const Order = std.math.Order; |
| 5 | |
| 6 | pub fn Treap(comptime Key: type, comptime compareFn: anytype) type { |
| 7 | return struct { |
| 8 | const Self = @This(); |
| 9 | |
| 10 | // Allow for compareFn to be fn(anytype, anytype) anytype |
| 11 | // which allows the convenient use of std.math.order. |
| 12 | fn compare(a: Key, b: Key) Order { |
| 13 | return compareFn(a, b); |
| 14 | } |
| 15 | |
| 16 | root: ?*Node = null, |
| 17 | prng: Prng = .{}, |
| 18 | |
| 19 | /// A customized pseudo random number generator for the treap. |
| 20 | /// This just helps reducing the memory size of the treap itself |
| 21 | /// as std.rand.DefaultPrng requires larger state (while producing better entropy for randomness to be fair). |
| 22 | const Prng = struct { |
| 23 | xorshift: usize = 0, |
| 24 | |
| 25 | fn random(self: *Prng, seed: usize) usize { |
| 26 | // Lazily seed the prng state |
| 27 | if (self.xorshift == 0) { |
| 28 | self.xorshift = seed; |
| 29 | } |
| 30 | |
| 31 | // Since we're using usize, decide the shifts by the integer's bit width. |
| 32 | const shifts = switch (@bitSizeOf(usize)) { |
| 33 | 64 => .{13, 7, 17}, |
| 34 | 32 => .{13, 17, 5}, |
| 35 | 16 => .{7, 9, 8}, |
| 36 | else => @compileError("platform not supported"), |
| 37 | }; |
| 38 | |
| 39 | self.xorshift ^= self.xorshift >> shifts[0]; |
| 40 | self.xorshift ^= self.xorshift << shifts[1]; |
| 41 | self.xorshift ^= self.xorshift >> shifts[2]; |
| 42 | |
| 43 | assert(self.xorshift != 0); |
| 44 | return self.xorshift; |
| 45 | } |
| 46 | }; |
| 47 | |
| 48 | /// A Node represents an item or point in the treap with a uniquely associated key. |
| 49 | pub const Node = struct { |
| 50 | key: Key, |
| 51 | priority: usize, |
| 52 | parent: ?*Node, |
| 53 | children: [2]?*Node, |
| 54 | }; |
| 55 | |
| 56 | /// Returns the smallest Node by key in the treap if there is one. |
| 57 | /// Use `getEntryForExisting()` to replace/remove this Node from the treap. |
| 58 | pub fn getMin(self: Self) ?*Node { |
| 59 | var node = self.root; |
| 60 | while (node) |current| { |
| 61 | node = current.children[0] orelse break; |
| 62 | } |
| 63 | return node; |
| 64 | } |
| 65 | |
| 66 | /// Returns the largest Node by key in the treap if there is one. |
| 67 | /// Use `getEntryForExisting()` to replace/remove this Node from the treap. |
| 68 | pub fn getMax(self: Self) ?*Node { |
| 69 | var node = self.root; |
| 70 | while (node) |current| { |
| 71 | node = current.children[1] orelse break; |
| 72 | } |
| 73 | return node; |
| 74 | } |
| 75 | |
| 76 | /// Lookup the Entry for the given key in the treap. |
| 77 | /// The Entry act's as a slot in the treap to insert/replace/remove the node associated with the key. |
| 78 | pub fn getEntryFor(self: *Self, key: Key) Entry { |
| 79 | var parent: ?*Node = undefined; |
| 80 | const node = self.find(key, &parent); |
| 81 | |
| 82 | return Entry{ |
| 83 | .key = key, |
| 84 | .treap = self, |
| 85 | .node = node, |
| 86 | .context = .{ .inserted_under = parent }, |
| 87 | }; |
| 88 | } |
| 89 | |
| 90 | /// Get an entry for a Node that currently exists in the treap. |
| 91 | /// It is undefined behavior if the Node is not currently inserted in the treap. |
| 92 | /// The Entry act's as a slot in the treap to insert/replace/remove the node associated with the key. |
| 93 | pub fn getEntryForExisting(self: *Self, node: *Node) Entry { |
| 94 | assert(node.priority != 0); |
| 95 | |
| 96 | return Entry{ |
| 97 | .key = node.key, |
| 98 | .treap = self, |
| 99 | .node = node, |
| 100 | .context = .{ .inserted_under = node.parent }, |
| 101 | }; |
| 102 | } |
| 103 | |
| 104 | /// An Entry represents a slot in the treap associated with a given key. |
| 105 | pub const Entry = struct { |
| 106 | key: Key, |
| 107 | treap: *Self, |
| 108 | node: ?*Node, |
| 109 | context: union(enum) { |
| 110 | /// A find() was called for this entry and the position in the treap is known. |
| 111 | inserted_under: ?*Node, |
| 112 | /// The entry's node was removed from the treap and a lookup must occur again for modification. |
| 113 | removed, |
| 114 | }, |
| 115 | |
| 116 | /// Returns the current Node at this Entry in the treap if there is one. |
| 117 | pub fn get(self: Entry) ?*Node { |
| 118 | return self.node; |
| 119 | } |
| 120 | |
| 121 | /// Update's the Node at this Entry in the treap with the new node. |
| 122 | pub fn set(self: *Entry, new_node: ?*Node) void { |
| 123 | // Update the entry's node reference after updating the treap below. |
| 124 | defer self.node = new_node; |
| 125 | |
| 126 | if (self.node) |old| { |
| 127 | if (new_node) |new| { |
| 128 | self.treap.replace(old, new); |
| 129 | return; |
| 130 | } |
| 131 | |
| 132 | self.treap.remove(old); |
| 133 | self.context = .removed; |
| 134 | return; |
| 135 | } |
| 136 | |
| 137 | if (new_node) |new| { |
| 138 | // A previous treap.remove() could have rebalanced the nodes |
| 139 | // so when inserting after a removal, we have to re-lookup the parent again. |
| 140 | // This lookup shouldn't find a node because we're yet to insert it.. |
| 141 | var parent: ?*Node = undefined; |
| 142 | switch (self.context) { |
| 143 | .inserted_under => |p| parent = p, |
| 144 | .removed => assert(self.treap.find(self.key, &parent) == null), |
| 145 | } |
| 146 | |
| 147 | self.treap.insert(self.key, parent, new); |
| 148 | self.context = .{ .inserted_under = parent }; |
| 149 | } |
| 150 | } |
| 151 | }; |
| 152 | |
| 153 | fn find(self: Self, key: Key, parent_ref: *?*Node) ?*Node { |
| 154 | var node = self.root; |
| 155 | parent_ref.* = null; |
| 156 | |
| 157 | // basic binary search while tracking the parent. |
| 158 | while (node) |current| { |
| 159 | const order = compare(key, current.key); |
| 160 | if (order == .eq) break; |
| 161 | |
| 162 | parent_ref.* = current; |
| 163 | node = current.children[@boolToInt(order == .gt)]; |
| 164 | } |
| 165 | |
| 166 | return node; |
| 167 | } |
| 168 | |
| 169 | fn insert(self: *Self, key: Key, parent: ?*Node, node: *Node) void { |
| 170 | // generate a random priority & prepare the node to be inserted into the tree |
| 171 | node.key = key; |
| 172 | node.priority = self.prng.random(@ptrToInt(node)); |
| 173 | node.parent = parent; |
| 174 | node.children = [_]?*Node{ null, null }; |
| 175 | |
| 176 | // point the parent at the new node |
| 177 | const link = if (parent) |p| &p.children[@boolToInt(compare(key, p.key) == .gt)] else &self.root; |
| 178 | assert(link.* == null); |
| 179 | link.* = node; |
| 180 | |
| 181 | // rotate the node up into the tree to balance it according to its priority |
| 182 | while (node.parent) |p| { |
| 183 | if (p.priority <= node.priority) break; |
| 184 | |
| 185 | const is_right = p.children[1] == @as(?*Node, node); |
| 186 | assert(p.children[@boolToInt(is_right)] == node); |
| 187 | |
| 188 | const rotate_right = !is_right; |
| 189 | self.rotate(p, rotate_right); |
| 190 | } |
| 191 | } |
| 192 | |
| 193 | fn replace(self: *Self, old: *Node, new: *Node) void { |
| 194 | // copy over the values from the old node |
| 195 | new.key = old.key; |
| 196 | new.priority = old.priority; |
| 197 | new.parent = old.parent; |
| 198 | new.children = old.children; |
| 199 | |
| 200 | // point the parent at the new node |
| 201 | const link = if (old.parent) |p| &p.children[@boolToInt(p.children[1] == old)] else &self.root; |
| 202 | assert(link.* == old); |
| 203 | link.* = new; |
| 204 | |
| 205 | // point the children's parent at the new node |
| 206 | for (old.children) |child_node| { |
| 207 | const child = child_node orelse continue; |
| 208 | assert(child.parent == old); |
| 209 | child.parent = new; |
| 210 | } |
| 211 | } |
| 212 | |
| 213 | fn remove(self: *Self, node: *Node) void { |
| 214 | // rotate the node down to be a leaf of the tree for removal, respecting priorities. |
| 215 | while (node.children[0] orelse node.children[1]) |_| { |
| 216 | self.rotate(node, rotate_right: { |
| 217 | const right = node.children[0] orelse break :rotate_right true; |
| 218 | const left = node.children[1] orelse break :rotate_right false; |
| 219 | break :rotate_right (left.priority < right.priority); |
| 220 | }); |
| 221 | } |
| 222 | |
| 223 | // node is a now a leaf; remove by nulling out the parent's reference to it. |
| 224 | const link = if (node.parent) |p| &p.children[@boolToInt(p.children[1] == node)] else &self.root; |
| 225 | assert(link.* == node); |
| 226 | link.* = null; |
| 227 | |
| 228 | // clean up after ourselves |
| 229 | node.key = undefined; |
| 230 | node.priority = 0; |
| 231 | node.parent = null; |
| 232 | node.children = [_]?*Node{ null, null }; |
| 233 | } |
| 234 | |
| 235 | fn rotate(self: *Self, node: *Node, right: bool) void { |
| 236 | // if right, converts the following: |
| 237 | // parent -> (node (target YY adjacent) XX) |
| 238 | // parent -> (target YY (node adjacent XX)) |
| 239 | // |
| 240 | // if left (!right), converts the following: |
| 241 | // parent -> (node (target YY adjacent) XX) |
| 242 | // parent -> (target YY (node adjacent XX)) |
| 243 | const parent = node.parent; |
| 244 | const target = node.children[@boolToInt(!right)] orelse unreachable; |
| 245 | const adjacent = target.children[@boolToInt(right)]; |
| 246 | |
| 247 | // do the rotation |
| 248 | target.children[@boolToInt(right)] = node; |
| 249 | node.parent = target; |
| 250 | node.children[@boolToInt(!right)] = adjacent; |
| 251 | if (adjacent) |adj| adj.parent = node; |
| 252 | |
| 253 | // fix the parent link |
| 254 | const link = if (parent) |p| &p.children[@boolToInt(p.children[1] == node)] else &self.root; |
| 255 | assert(link.* == node); |
| 256 | link.* = target; |
| 257 | } |
| 258 | }; |
| 259 | } |
| 260 | |
| 261 | const TestTreap = Treap(u64, std.math.order); |
| 262 | const TestNode = TestTreap.Node; |
| 263 | |
| 264 | test "std.Treap: insert, find, remove" { |
| 265 | var prng = std.rand.DefaultPrng.init(0xdeadbeef); |
| 266 | var rng = prng.random(); |
| 267 | |
| 268 | var treap = TestTreap{}; |
| 269 | var nodes: [6]TestNode = undefined; |
| 270 | |
| 271 | for (nodes) |*node| { |
| 272 | const key = rng.int(u64); |
| 273 | |
| 274 | var entry = treap.getEntryFor(key); |
| 275 | try testing.expectEqual(entry.key, key); |
| 276 | try testing.expectEqual(entry.get(), null); |
| 277 | |
| 278 | entry.set(node); |
| 279 | try testing.expectEqual(entry.key, key); |
| 280 | try testing.expectEqual(node.key, key); |
| 281 | try testing.expectEqual(entry.get(), node); |
| 282 | } |
| 283 | |
| 284 | for (nodes) |*node| { |
| 285 | const key = node.key; |
| 286 | |
| 287 | var entry = treap.getEntryFor(node.key); |
| 288 | try testing.expectEqual(entry.key, key); |
| 289 | try testing.expectEqual(entry.get(), node); |
| 290 | |
| 291 | var existingEntry = treap.getEntryForExisting(node); |
| 292 | try testing.expectEqual(entry, existingEntry); |
| 293 | } |
| 294 | } |
kprotty <kbutcher6200@gmail.com> 2022-04-15 15:26:51-05:00