1const std = @import("std.zig");
2const assert = std.debug.assert;
3const testing = std.testing;
4const Order = std.math.Order;
5
6pub 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.Random.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 pub fn next(node: *Node) ?*Node {
56 return nextOnDirection(node, 1);
57 }
58 pub fn prev(node: *Node) ?*Node {
59 return nextOnDirection(node, 0);
60 }
61 };
62
63 fn extremeInSubtreeOnDirection(node: *Node, direction: u1) *Node {
64 var cur = node;
65 while (cur.children[direction]) |next| cur = next;
66 return cur;
67 }
68
69 fn nextOnDirection(node: *Node, direction: u1) ?*Node {
70 if (node.children[direction]) |child| {
71 return extremeInSubtreeOnDirection(child, direction ^ 1);
72 }
73 var cur = node;
74 // Traversing upward until we find `parent` to `cur` is NOT on
75 // `direction`, or equivalently, `cur` to `parent` IS on
76 // `direction` thus `parent` is the next.
77 while (true) {
78 if (cur.parent) |parent| {
79 // If `parent -> node` is NOT on `direction`, then
80 // `node -> parent` IS on `direction`
81 if (parent.children[direction] != cur) return parent;
82 cur = parent;
83 } else {
84 return null;
85 }
86 }
87 }
88
89 /// Returns the smallest Node by key in the treap if there is one.
90 /// Use `getEntryForExisting()` to replace/remove this Node from the treap.
91 pub fn getMin(self: Self) ?*Node {
92 if (self.root) |root| return extremeInSubtreeOnDirection(root, 0);
93 return null;
94 }
95
96 /// Returns the largest Node by key in the treap if there is one.
97 /// Use `getEntryForExisting()` to replace/remove this Node from the treap.
98 pub fn getMax(self: Self) ?*Node {
99 if (self.root) |root| return extremeInSubtreeOnDirection(root, 1);
100 return null;
101 }
102
103 /// Lookup the Entry for the given key in the treap.
104 /// The Entry act's as a slot in the treap to insert/replace/remove the node associated with the key.
105 pub fn getEntryFor(self: *Self, key: Key) Entry {
106 var parent: ?*Node = undefined;
107 const node = self.find(key, &parent);
108
109 return Entry{
110 .key = key,
111 .treap = self,
112 .node = node,
113 .context = .{ .inserted_under = parent },
114 };
115 }
116
117 /// Get an entry for a Node that currently exists in the treap.
118 /// It is undefined behavior if the Node is not currently inserted in the treap.
119 /// The Entry act's as a slot in the treap to insert/replace/remove the node associated with the key.
120 pub fn getEntryForExisting(self: *Self, node: *Node) Entry {
121 assert(node.priority != 0);
122
123 return Entry{
124 .key = node.key,
125 .treap = self,
126 .node = node,
127 .context = .{ .inserted_under = node.parent },
128 };
129 }
130
131 /// An Entry represents a slot in the treap associated with a given key.
132 pub const Entry = struct {
133 /// The associated key for this entry.
134 key: Key,
135 /// A reference to the treap this entry is apart of.
136 treap: *Self,
137 /// The current node at this entry.
138 node: ?*Node,
139 /// The current state of the entry.
140 context: union(enum) {
141 /// A find() was called for this entry and the position in the treap is known.
142 inserted_under: ?*Node,
143 /// The entry's node was removed from the treap and a lookup must occur again for modification.
144 removed,
145 },
146
147 /// Update's the Node at this Entry in the treap with the new node (null for deleting). `new_node`
148 /// can have `undefind` content because the value will be initialized internally.
149 pub fn set(self: *Entry, new_node: ?*Node) void {
150 // Update the entry's node reference after updating the treap below.
151 defer self.node = new_node;
152
153 if (self.node) |old| {
154 if (new_node) |new| {
155 self.treap.replace(old, new);
156 return;
157 }
158
159 self.treap.remove(old);
160 self.context = .removed;
161 return;
162 }
163
164 if (new_node) |new| {
165 // A previous treap.remove() could have rebalanced the nodes
166 // so when inserting after a removal, we have to re-lookup the parent again.
167 // This lookup shouldn't find a node because we're yet to insert it..
168 var parent: ?*Node = undefined;
169 switch (self.context) {
170 .inserted_under => |p| parent = p,
171 .removed => assert(self.treap.find(self.key, &parent) == null),
172 }
173
174 self.treap.insert(self.key, parent, new);
175 self.context = .{ .inserted_under = parent };
176 }
177 }
178 };
179
180 fn find(self: Self, key: Key, parent_ref: *?*Node) ?*Node {
181 var node = self.root;
182 parent_ref.* = null;
183
184 // basic binary search while tracking the parent.
185 while (node) |current| {
186 const order = compare(key, current.key);
187 if (order == .eq) break;
188
189 parent_ref.* = current;
190 node = current.children[@intFromBool(order == .gt)];
191 }
192
193 return node;
194 }
195
196 fn insert(self: *Self, key: Key, parent: ?*Node, node: *Node) void {
197 // generate a random priority & prepare the node to be inserted into the tree
198 node.key = key;
199 node.priority = self.prng.random(@intFromPtr(node));
200 node.parent = parent;
201 node.children = [_]?*Node{ null, null };
202
203 // point the parent at the new node
204 const link = if (parent) |p| &p.children[@intFromBool(compare(key, p.key) == .gt)] else &self.root;
205 assert(link.* == null);
206 link.* = node;
207
208 // rotate the node up into the tree to balance it according to its priority
209 while (node.parent) |p| {
210 if (p.priority <= node.priority) break;
211
212 const is_right = p.children[1] == node;
213 assert(p.children[@intFromBool(is_right)] == node);
214
215 const rotate_right = !is_right;
216 self.rotate(p, rotate_right);
217 }
218 }
219
220 fn replace(self: *Self, old: *Node, new: *Node) void {
221 // copy over the values from the old node
222 new.key = old.key;
223 new.priority = old.priority;
224 new.parent = old.parent;
225 new.children = old.children;
226
227 // point the parent at the new node
228 const link = if (old.parent) |p| &p.children[@intFromBool(p.children[1] == old)] else &self.root;
229 assert(link.* == old);
230 link.* = new;
231
232 // point the children's parent at the new node
233 for (old.children) |child_node| {
234 const child = child_node orelse continue;
235 assert(child.parent == old);
236 child.parent = new;
237 }
238 }
239
240 fn remove(self: *Self, node: *Node) void {
241 // rotate the node down to be a leaf of the tree for removal, respecting priorities.
242 while (node.children[0] orelse node.children[1]) |_| {
243 self.rotate(node, rotate_right: {
244 const right = node.children[1] orelse break :rotate_right true;
245 const left = node.children[0] orelse break :rotate_right false;
246 break :rotate_right (left.priority < right.priority);
247 });
248 }
249
250 // node is a now a leaf; remove by nulling out the parent's reference to it.
251 const link = if (node.parent) |p| &p.children[@intFromBool(p.children[1] == node)] else &self.root;
252 assert(link.* == node);
253 link.* = null;
254
255 // clean up after ourselves
256 node.priority = 0;
257 node.parent = null;
258 node.children = [_]?*Node{ null, null };
259 }
260
261 fn rotate(self: *Self, node: *Node, right: bool) void {
262 // if right, converts the following:
263 // parent -> (node (target YY adjacent) XX)
264 // parent -> (target YY (node adjacent XX))
265 //
266 // if left (!right), converts the following:
267 // parent -> (node (target YY adjacent) XX)
268 // parent -> (target YY (node adjacent XX))
269 const parent = node.parent;
270 const target = node.children[@intFromBool(!right)] orelse unreachable;
271 const adjacent = target.children[@intFromBool(right)];
272
273 // rotate the children
274 target.children[@intFromBool(right)] = node;
275 node.children[@intFromBool(!right)] = adjacent;
276
277 // rotate the parents
278 node.parent = target;
279 target.parent = parent;
280 if (adjacent) |adj| adj.parent = node;
281
282 // fix the parent link
283 const link = if (parent) |p| &p.children[@intFromBool(p.children[1] == node)] else &self.root;
284 assert(link.* == node);
285 link.* = target;
286 }
287
288 /// Usage example:
289 /// var iter = treap.inorderIterator();
290 /// while (iter.next()) |node| {
291 /// ...
292 /// }
293 pub const InorderIterator = struct {
294 current: ?*Node,
295
296 pub fn next(it: *InorderIterator) ?*Node {
297 const current = it.current;
298 it.current = if (current) |cur|
299 cur.next()
300 else
301 null;
302 return current;
303 }
304 };
305
306 pub fn inorderIterator(self: *Self) InorderIterator {
307 return .{ .current = self.getMin() };
308 }
309 };
310}
311
312// For iterating a slice in a random order
313// https://lemire.me/blog/2017/09/18/visiting-all-values-in-an-array-exactly-once-in-random-order/
314fn SliceIterRandomOrder(comptime T: type) type {
315 return struct {
316 rng: std.Random,
317 slice: []T,
318 index: usize = undefined,
319 offset: usize = undefined,
320 co_prime: usize,
321
322 const Self = @This();
323
324 pub fn init(slice: []T, rng: std.Random) Self {
325 return Self{
326 .rng = rng,
327 .slice = slice,
328 .co_prime = blk: {
329 if (slice.len == 0) break :blk 0;
330 var prime = slice.len / 2;
331 while (prime < slice.len) : (prime += 1) {
332 var gcd = [_]usize{ prime, slice.len };
333 while (gcd[1] != 0) {
334 const temp = gcd;
335 gcd = [_]usize{ temp[1], temp[0] % temp[1] };
336 }
337 if (gcd[0] == 1) break;
338 }
339 break :blk prime;
340 },
341 };
342 }
343
344 pub fn reset(self: *Self) void {
345 self.index = 0;
346 self.offset = self.rng.int(usize);
347 }
348
349 pub fn next(self: *Self) ?*T {
350 if (self.index >= self.slice.len) return null;
351 defer self.index += 1;
352 return &self.slice[((self.index *% self.co_prime) +% self.offset) % self.slice.len];
353 }
354 };
355}
356
357const TestTreap = Treap(u64, std.math.order);
358const TestNode = TestTreap.Node;
359
360test "insert, find, replace, remove" {
361 var treap = TestTreap{};
362 var nodes: [10]TestNode = undefined;
363
364 var prng = std.Random.DefaultPrng.init(0xdeadbeef);
365 var iter = SliceIterRandomOrder(TestNode).init(&nodes, prng.random());
366
367 // insert check
368 iter.reset();
369 while (iter.next()) |node| {
370 const key = prng.random().int(u64);
371
372 // make sure the current entry is empty.
373 var entry = treap.getEntryFor(key);
374 try testing.expectEqual(entry.key, key);
375 try testing.expectEqual(entry.node, null);
376
377 // insert the entry and make sure the fields are correct.
378 entry.set(node);
379 try testing.expectEqual(node.key, key);
380 try testing.expectEqual(entry.key, key);
381 try testing.expectEqual(entry.node, node);
382 }
383
384 // find check
385 iter.reset();
386 while (iter.next()) |node| {
387 const key = node.key;
388
389 // find the entry by-key and by-node after having been inserted.
390 const entry = treap.getEntryFor(node.key);
391 try testing.expectEqual(entry.key, key);
392 try testing.expectEqual(entry.node, node);
393 try testing.expectEqual(entry.node, treap.getEntryForExisting(node).node);
394 }
395
396 // in-order iterator check
397 {
398 var it = treap.inorderIterator();
399 var last_key: u64 = 0;
400 while (it.next()) |node| {
401 try std.testing.expect(node.key >= last_key);
402 last_key = node.key;
403 }
404 }
405
406 // replace check
407 iter.reset();
408 while (iter.next()) |node| {
409 const key = node.key;
410
411 // find the entry by node since we already know it exists
412 var entry = treap.getEntryForExisting(node);
413 try testing.expectEqual(entry.key, key);
414 try testing.expectEqual(entry.node, node);
415
416 var stub_node: TestNode = undefined;
417
418 // replace the node with a stub_node and ensure future finds point to the stub_node.
419 entry.set(&stub_node);
420 try testing.expectEqual(entry.node, &stub_node);
421 try testing.expectEqual(entry.node, treap.getEntryFor(key).node);
422 try testing.expectEqual(entry.node, treap.getEntryForExisting(&stub_node).node);
423
424 // replace the stub_node back to the node and ensure future finds point to the old node.
425 entry.set(node);
426 try testing.expectEqual(entry.node, node);
427 try testing.expectEqual(entry.node, treap.getEntryFor(key).node);
428 try testing.expectEqual(entry.node, treap.getEntryForExisting(node).node);
429 }
430
431 // remove check
432 iter.reset();
433 while (iter.next()) |node| {
434 const key = node.key;
435
436 // find the entry by node since we already know it exists
437 var entry = treap.getEntryForExisting(node);
438 try testing.expectEqual(entry.key, key);
439 try testing.expectEqual(entry.node, node);
440
441 // remove the node at the entry and ensure future finds point to it being removed.
442 entry.set(null);
443 try testing.expectEqual(entry.node, null);
444 try testing.expectEqual(entry.node, treap.getEntryFor(key).node);
445
446 // insert the node back and ensure future finds point to the inserted node
447 entry.set(node);
448 try testing.expectEqual(entry.node, node);
449 try testing.expectEqual(entry.node, treap.getEntryFor(key).node);
450 try testing.expectEqual(entry.node, treap.getEntryForExisting(node).node);
451
452 // remove the node again and make sure it was cleared after the insert
453 entry.set(null);
454 try testing.expectEqual(entry.node, null);
455 try testing.expectEqual(entry.node, treap.getEntryFor(key).node);
456 }
457}
458
459test "inorderIterator" {
460 var treap = TestTreap{};
461 var nodes: [10]TestNode = undefined;
462
463 // Build the tree.
464 var i: usize = 0;
465 while (i < 10) : (i += 1) {
466 const key = @as(u64, i);
467 var entry = treap.getEntryFor(key);
468 entry.set(&nodes[i]);
469 }
470
471 // Test the iterator.
472 var iter = treap.inorderIterator();
473 i = 0;
474 while (iter.next()) |node| {
475 const key = @as(u64, i);
476 try testing.expectEqual(key, node.key);
477 i += 1;
478 }
479}
480
481test "getMin, getMax, simple" {
482 var treap = TestTreap{};
483 var nodes: [3]TestNode = undefined;
484
485 try testing.expectEqual(null, treap.getMin());
486 try testing.expectEqual(null, treap.getMax());
487 { // nodes[1]
488 var entry = treap.getEntryFor(1);
489 entry.set(&nodes[1]);
490 try testing.expectEqual(&nodes[1], treap.getMin());
491 try testing.expectEqual(&nodes[1], treap.getMax());
492 }
493 { // nodes[0]
494 var entry = treap.getEntryFor(0);
495 entry.set(&nodes[0]);
496 try testing.expectEqual(&nodes[0], treap.getMin());
497 try testing.expectEqual(&nodes[1], treap.getMax());
498 }
499 { // nodes[2]
500 var entry = treap.getEntryFor(2);
501 entry.set(&nodes[2]);
502 try testing.expectEqual(&nodes[0], treap.getMin());
503 try testing.expectEqual(&nodes[2], treap.getMax());
504 }
505}
506
507test "getMin, getMax, random" {
508 var nodes: [100]TestNode = undefined;
509 var prng = std.Random.DefaultPrng.init(0xdeadbeef);
510 var iter = SliceIterRandomOrder(TestNode).init(&nodes, prng.random());
511
512 var treap = TestTreap{};
513 var min: u64 = std.math.maxInt(u64);
514 var max: u64 = 0;
515
516 try testing.expectEqual(null, treap.getMin());
517 try testing.expectEqual(null, treap.getMax());
518
519 // Insert and check min/max after each insertion.
520 iter.reset();
521 while (iter.next()) |node| {
522 const key = prng.random().int(u64);
523
524 // Insert into `treap`.
525 var entry = treap.getEntryFor(key);
526 entry.set(node);
527
528 if (key < min) min = key;
529 if (key > max) max = key;
530
531 const min_node = treap.getMin().?;
532 try std.testing.expectEqual(null, min_node.prev());
533 try std.testing.expectEqual(min, min_node.key);
534
535 const max_node = treap.getMax().?;
536 try std.testing.expectEqual(null, max_node.next());
537 try std.testing.expectEqual(max, max_node.key);
538 }
539}
540
541test "node.{prev(),next()} with sequential insertion and deletion" {
542 // Insert order: 50, 0, 1, 2, ..., 49, 51, 52, ..., 99.
543 // Delete order: 0, 1, 2, ..., 49, 51, 52, ..., 99.
544 // Check 50's neighbors.
545 var treap = TestTreap{};
546 var nodes: [100]TestNode = undefined;
547 {
548 var entry = treap.getEntryFor(50);
549 entry.set(&nodes[50]);
550 try testing.expectEqual(50, nodes[50].key);
551 try testing.expectEqual(null, nodes[50].prev());
552 try testing.expectEqual(null, nodes[50].next());
553 }
554 // Insert others.
555 var i: usize = 0;
556 while (i < 50) : (i += 1) {
557 const key = @as(u64, i);
558 const node = &nodes[i];
559 var entry = treap.getEntryFor(key);
560 entry.set(node);
561 try testing.expectEqual(key, node.key);
562 try testing.expectEqual(node, nodes[50].prev());
563 try testing.expectEqual(null, nodes[50].next());
564 }
565 i = 51;
566 while (i < 100) : (i += 1) {
567 const key = @as(u64, i);
568 const node = &nodes[i];
569 var entry = treap.getEntryFor(key);
570 entry.set(node);
571 try testing.expectEqual(key, node.key);
572 try testing.expectEqual(&nodes[49], nodes[50].prev());
573 try testing.expectEqual(&nodes[51], nodes[50].next());
574 }
575 // Remove others.
576 i = 0;
577 while (i < 49) : (i += 1) {
578 const key = @as(u64, i);
579 var entry = treap.getEntryFor(key);
580 entry.set(null);
581 try testing.expectEqual(&nodes[49], nodes[50].prev());
582 try testing.expectEqual(&nodes[51], nodes[50].next());
583 }
584 { // i = 49.
585 const key = @as(u64, i);
586 var entry = treap.getEntryFor(key);
587 entry.set(null);
588 try testing.expectEqual(null, nodes[50].prev());
589 try testing.expectEqual(&nodes[51], nodes[50].next());
590 }
591 i = 51;
592 while (i < 99) : (i += 1) {
593 const key = @as(u64, i);
594 var entry = treap.getEntryFor(key);
595 entry.set(null);
596 try testing.expectEqual(null, nodes[50].prev());
597 try testing.expectEqual(&nodes[i + 1], nodes[50].next());
598 }
599 { // i = 99.
600 const key = @as(u64, i);
601 var entry = treap.getEntryFor(key);
602 entry.set(null);
603 try testing.expectEqual(null, nodes[50].prev());
604 try testing.expectEqual(null, nodes[50].next());
605 }
606}
607
608fn findFirstGreaterOrEqual(array: []u64, value: u64) usize {
609 var i: usize = 0;
610 while (i < array.len and array[i] < value) i += 1;
611 return i;
612}
613
614fn testOrderedArrayAndTreapConsistency(array: []u64, treap: *TestTreap) !void {
615 var i: usize = 0;
616 while (i < array.len) : (i += 1) {
617 const value = array[i];
618
619 const entry = treap.getEntryFor(value);
620 try testing.expect(entry.node != null);
621 const node = entry.node.?;
622 try testing.expectEqual(value, node.key);
623
624 if (i == 0) {
625 try testing.expectEqual(node.prev(), null);
626 } else {
627 try testing.expectEqual(node.prev(), treap.getEntryFor(array[i - 1]).node);
628 }
629 if (i + 1 == array.len) {
630 try testing.expectEqual(node.next(), null);
631 } else {
632 try testing.expectEqual(node.next(), treap.getEntryFor(array[i + 1]).node);
633 }
634 }
635}
636
637test "node.{prev(),next()} with random data" {
638 var nodes: [100]TestNode = undefined;
639 var prng = std.Random.DefaultPrng.init(0xdeadbeef);
640 var iter = SliceIterRandomOrder(TestNode).init(&nodes, prng.random());
641
642 var treap = TestTreap{};
643 // A slow, stupid but correct reference. Ordered.
644 var golden = std.array_list.Managed(u64).init(std.testing.allocator);
645 defer golden.deinit();
646
647 // Insert.
648 iter.reset();
649 while (iter.next()) |node| {
650 const key = prng.random().int(u64);
651
652 // Insert into `golden`.
653 const i = findFirstGreaterOrEqual(golden.items, key);
654 // Ensure not found. If found: `prng`'s fault.
655 try testing.expect(i == golden.items.len or golden.items[i] > key);
656 try golden.insert(i, key);
657
658 // Insert into `treap`.
659 var entry = treap.getEntryFor(key);
660 entry.set(node);
661
662 try testOrderedArrayAndTreapConsistency(golden.items, &treap);
663 }
664
665 // Delete.
666 iter.reset();
667 while (iter.next()) |node| {
668 const key = node.key;
669
670 // Delete from `golden`.
671 const i = findFirstGreaterOrEqual(golden.items, key);
672 try testing.expect(i < golden.items.len);
673 _ = golden.orderedRemove(i);
674
675 // Delete from `treap`.
676 var entry = treap.getEntryFor(key);
677 try testing.expect(entry.node != null);
678 entry.set(null);
679
680 try testOrderedArrayAndTreapConsistency(golden.items, &treap);
681 }
682}