| ... | @@ -0,0 +1,398 @@ |
| 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 | /// The associated key for this entry. |
| 107 | key: Key, |
| 108 | /// A reference to the treap this entry is apart of. |
| 109 | treap: *Self, |
| 110 | /// The current node at this entry. |
| 111 | node: ?*Node, |
| 112 | /// The current state of the entry. |
| 113 | context: union(enum) { |
| 114 | /// A find() was called for this entry and the position in the treap is known. |
| 115 | inserted_under: ?*Node, |
| 116 | /// The entry's node was removed from the treap and a lookup must occur again for modification. |
| 117 | removed, |
| 118 | }, |
| 119 | |
| 120 | /// Update's the Node at this Entry in the treap with the new node. |
| 121 | pub fn set(self: *Entry, new_node: ?*Node) void { |
| 122 | // Update the entry's node reference after updating the treap below. |
| 123 | defer self.node = new_node; |
| 124 | |
| 125 | if (self.node) |old| { |
| 126 | if (new_node) |new| { |
| 127 | self.treap.replace(old, new); |
| 128 | return; |
| 129 | } |
| 130 | |
| 131 | self.treap.remove(old); |
| 132 | self.context = .removed; |
| 133 | return; |
| 134 | } |
| 135 | |
| 136 | if (new_node) |new| { |
| 137 | // A previous treap.remove() could have rebalanced the nodes |
| 138 | // so when inserting after a removal, we have to re-lookup the parent again. |
| 139 | // This lookup shouldn't find a node because we're yet to insert it.. |
| 140 | var parent: ?*Node = undefined; |
| 141 | switch (self.context) { |
| 142 | .inserted_under => |p| parent = p, |
| 143 | .removed => assert(self.treap.find(self.key, &parent) == null), |
| 144 | } |
| 145 | |
| 146 | self.treap.insert(self.key, parent, new); |
| 147 | self.context = .{ .inserted_under = parent }; |
| 148 | } |
| 149 | } |
| 150 | }; |
| 151 | |
| 152 | fn find(self: Self, key: Key, parent_ref: *?*Node) ?*Node { |
| 153 | var node = self.root; |
| 154 | parent_ref.* = null; |
| 155 | |
| 156 | // basic binary search while tracking the parent. |
| 157 | while (node) |current| { |
| 158 | const order = compare(key, current.key); |
| 159 | if (order == .eq) break; |
| 160 | |
| 161 | parent_ref.* = current; |
| 162 | node = current.children[@boolToInt(order == .gt)]; |
| 163 | } |
| 164 | |
| 165 | return node; |
| 166 | } |
| 167 | |
| 168 | fn insert(self: *Self, key: Key, parent: ?*Node, node: *Node) void { |
| 169 | // generate a random priority & prepare the node to be inserted into the tree |
| 170 | node.key = key; |
| 171 | node.priority = self.prng.random(@ptrToInt(node)); |
| 172 | node.parent = parent; |
| 173 | node.children = [_]?*Node{ null, null }; |
| 174 | |
| 175 | // point the parent at the new node |
| 176 | const link = if (parent) |p| &p.children[@boolToInt(compare(key, p.key) == .gt)] else &self.root; |
| 177 | assert(link.* == null); |
| 178 | link.* = node; |
| 179 | |
| 180 | // rotate the node up into the tree to balance it according to its priority |
| 181 | while (node.parent) |p| { |
| 182 | if (p.priority <= node.priority) break; |
| 183 | |
| 184 | const is_right = p.children[1] == node; |
| 185 | assert(p.children[@boolToInt(is_right)] == node); |
| 186 | |
| 187 | const rotate_right = !is_right; |
| 188 | self.rotate(p, rotate_right); |
| 189 | } |
| 190 | } |
| 191 | |
| 192 | fn replace(self: *Self, old: *Node, new: *Node) void { |
| 193 | // copy over the values from the old node |
| 194 | new.key = old.key; |
| 195 | new.priority = old.priority; |
| 196 | new.parent = old.parent; |
| 197 | new.children = old.children; |
| 198 | |
| 199 | // point the parent at the new node |
| 200 | const link = if (old.parent) |p| &p.children[@boolToInt(p.children[1] == old)] else &self.root; |
| 201 | assert(link.* == old); |
| 202 | link.* = new; |
| 203 | |
| 204 | // point the children's parent at the new node |
| 205 | for (old.children) |child_node| { |
| 206 | const child = child_node orelse continue; |
| 207 | assert(child.parent == old); |
| 208 | child.parent = new; |
| 209 | } |
| 210 | } |
| 211 | |
| 212 | fn remove(self: *Self, node: *Node) void { |
| 213 | // rotate the node down to be a leaf of the tree for removal, respecting priorities. |
| 214 | while (node.children[0] orelse node.children[1]) |_| { |
| 215 | self.rotate(node, rotate_right: { |
| 216 | const right = node.children[1] orelse break :rotate_right true; |
| 217 | const left = node.children[0] orelse break :rotate_right false; |
| 218 | break :rotate_right (left.priority < right.priority); |
| 219 | }); |
| 220 | } |
| 221 | |
| 222 | // node is a now a leaf; remove by nulling out the parent's reference to it. |
| 223 | const link = if (node.parent) |p| &p.children[@boolToInt(p.children[1] == node)] else &self.root; |
| 224 | assert(link.* == node); |
| 225 | link.* = null; |
| 226 | |
| 227 | // clean up after ourselves |
| 228 | node.key = undefined; |
| 229 | node.priority = 0; |
| 230 | node.parent = null; |
| 231 | node.children = [_]?*Node{ null, null }; |
| 232 | } |
| 233 | |
| 234 | fn rotate(self: *Self, node: *Node, right: bool) void { |
| 235 | // if right, converts the following: |
| 236 | // parent -> (node (target YY adjacent) XX) |
| 237 | // parent -> (target YY (node adjacent XX)) |
| 238 | // |
| 239 | // if left (!right), converts the following: |
| 240 | // parent -> (node (target YY adjacent) XX) |
| 241 | // parent -> (target YY (node adjacent XX)) |
| 242 | const parent = node.parent; |
| 243 | const target = node.children[@boolToInt(!right)] orelse unreachable; |
| 244 | const adjacent = target.children[@boolToInt(right)]; |
| 245 | |
| 246 | // rotate the children |
| 247 | target.children[@boolToInt(right)] = node; |
| 248 | node.children[@boolToInt(!right)] = adjacent; |
| 249 | |
| 250 | // rotate the parents |
| 251 | node.parent = target; |
| 252 | target.parent = parent; |
| 253 | if (adjacent) |adj| adj.parent = node; |
| 254 | |
| 255 | // fix the parent link |
| 256 | const link = if (parent) |p| &p.children[@boolToInt(p.children[1] == node)] else &self.root; |
| 257 | assert(link.* == node); |
| 258 | link.* = target; |
| 259 | } |
| 260 | }; |
| 261 | } |
| 262 | |
| 263 | // For iterating a slice in a random order |
| 264 | // https://lemire.me/blog/2017/09/18/visiting-all-values-in-an-array-exactly-once-in-random-order/ |
| 265 | fn SliceIterRandomOrder(comptime T: type) type { |
| 266 | return struct { |
| 267 | rng: std.rand.Random, |
| 268 | slice: []T, |
| 269 | index: usize = undefined, |
| 270 | offset: usize = undefined, |
| 271 | co_prime: usize, |
| 272 | |
| 273 | const Self = @This(); |
| 274 | |
| 275 | pub fn init(slice: []T, rng: std.rand.Random) Self { |
| 276 | return Self{ |
| 277 | .rng = rng, |
| 278 | .slice = slice, |
| 279 | .co_prime = blk: { |
| 280 | if (slice.len == 0) break :blk 0; |
| 281 | var prime = slice.len / 2; |
| 282 | while (prime < slice.len) : (prime += 1) { |
| 283 | var gcd = [_]usize{ prime, slice.len }; |
| 284 | while (gcd[1] != 0) { |
| 285 | const temp = gcd; |
| 286 | gcd = [_]usize{ temp[1], temp[0] % temp[1] }; |
| 287 | } |
| 288 | if (gcd[0] == 1) break; |
| 289 | } |
| 290 | break :blk prime; |
| 291 | }, |
| 292 | }; |
| 293 | } |
| 294 | |
| 295 | pub fn reset(self: *Self) void { |
| 296 | self.index = 0; |
| 297 | self.offset = self.rng.int(usize); |
| 298 | } |
| 299 | |
| 300 | pub fn next(self: *Self) ?*T { |
| 301 | if (self.index >= self.slice.len) return null; |
| 302 | defer self.index += 1; |
| 303 | return &self.slice[((self.index *% self.co_prime) +% self.offset) % self.slice.len]; |
| 304 | } |
| 305 | }; |
| 306 | } |
| 307 | |
| 308 | const TestTreap = Treap(u64, std.math.order); |
| 309 | const TestNode = TestTreap.Node; |
| 310 | |
| 311 | test "std.Treap: insert, find, replace, remove" { |
| 312 | var treap = TestTreap{}; |
| 313 | var nodes: [10]TestNode = undefined; |
| 314 | |
| 315 | var prng = std.rand.DefaultPrng.init(0xdeadbeef); |
| 316 | var iter = SliceIterRandomOrder(TestNode).init(&nodes, prng.random()); |
| 317 | |
| 318 | // insert check |
| 319 | iter.reset(); |
| 320 | while (iter.next()) |node| { |
| 321 | const key = prng.random().int(u64); |
| 322 | |
| 323 | // make sure the current entry is empty. |
| 324 | var entry = treap.getEntryFor(key); |
| 325 | try testing.expectEqual(entry.key, key); |
| 326 | try testing.expectEqual(entry.node, null); |
| 327 | |
| 328 | // insert the entry and make sure the fields are correct. |
| 329 | entry.set(node); |
| 330 | try testing.expectEqual(node.key, key); |
| 331 | try testing.expectEqual(entry.key, key); |
| 332 | try testing.expectEqual(entry.node, node); |
| 333 | } |
| 334 | |
| 335 | // find check |
| 336 | iter.reset(); |
| 337 | while (iter.next()) |node| { |
| 338 | const key = node.key; |
| 339 | |
| 340 | // find the entry by-key and by-node after having been inserted. |
| 341 | var entry = treap.getEntryFor(node.key); |
| 342 | try testing.expectEqual(entry.key, key); |
| 343 | try testing.expectEqual(entry.node, node); |
| 344 | try testing.expectEqual(entry.node, treap.getEntryForExisting(node).node); |
| 345 | } |
| 346 | |
| 347 | // replace check |
| 348 | iter.reset(); |
| 349 | while (iter.next()) |node| { |
| 350 | const key = node.key; |
| 351 | |
| 352 | // find the entry by node since we already know it exists |
| 353 | var entry = treap.getEntryForExisting(node); |
| 354 | try testing.expectEqual(entry.key, key); |
| 355 | try testing.expectEqual(entry.node, node); |
| 356 | |
| 357 | var stub_node: TestNode = undefined; |
| 358 | |
| 359 | // replace the node with a stub_node and ensure future finds point to the stub_node. |
| 360 | entry.set(&stub_node); |
| 361 | try testing.expectEqual(entry.node, &stub_node); |
| 362 | try testing.expectEqual(entry.node, treap.getEntryFor(key).node); |
| 363 | try testing.expectEqual(entry.node, treap.getEntryForExisting(&stub_node).node); |
| 364 | |
| 365 | // replace the stub_node back to the node and ensure future finds point to the old node. |
| 366 | entry.set(node); |
| 367 | try testing.expectEqual(entry.node, node); |
| 368 | try testing.expectEqual(entry.node, treap.getEntryFor(key).node); |
| 369 | try testing.expectEqual(entry.node, treap.getEntryForExisting(node).node); |
| 370 | } |
| 371 | |
| 372 | // remove check |
| 373 | iter.reset(); |
| 374 | while (iter.next()) |node| { |
| 375 | const key = node.key; |
| 376 | |
| 377 | // find the entry by node since we already know it exists |
| 378 | var entry = treap.getEntryForExisting(node); |
| 379 | try testing.expectEqual(entry.key, key); |
| 380 | try testing.expectEqual(entry.node, node); |
| 381 | |
| 382 | // remove the node at the entry and ensure future finds point to it being removed. |
| 383 | entry.set(null); |
| 384 | try testing.expectEqual(entry.node, null); |
| 385 | try testing.expectEqual(entry.node, treap.getEntryFor(key).node); |
| 386 | |
| 387 | // insert the node back and ensure future finds point to the inserted node |
| 388 | entry.set(node); |
| 389 | try testing.expectEqual(entry.node, node); |
| 390 | try testing.expectEqual(entry.node, treap.getEntryFor(key).node); |
| 391 | try testing.expectEqual(entry.node, treap.getEntryForExisting(node).node); |
| 392 | |
| 393 | // remove the node again and make sure it was cleared after the insert |
| 394 | entry.set(null); |
| 395 | try testing.expectEqual(entry.node, null); |
| 396 | try testing.expectEqual(entry.node, treap.getEntryFor(key).node); |
| 397 | } |
| 398 | } |
Andrew Kelley <andrew@ziglang.org> 2022-04-16 06:23:18-04:00
GitHub <noreply@github.com> 2022-04-16 06:23:18-04:00
Signed by PGP key 4AEE18F83AFDEB23