authorgravatar for git@zander.xyzZander Khan <git@zander.xyz> 2021-01-16 12:01:06+00:00
committergravatar for git@zander.xyzZander Khan <git@zander.xyz> 2021-01-16 12:01:06+00:00
loga727a508cd18d4af1773841e00630f5b47e2b95f
tree5117c85f8252cda98ef9d5d703fe9dced7012979
parentb204ea0349d5a580fc8ba9d8059c520301072072

std: Add Priority Dequeue

2 files changed, 880 insertions(+), 0 deletions(-)

lib/std/priority_dequeue.zig created+879
......@@ -0,0 +1,879 @@
1const std = @import("std");
2const Allocator = std.mem.Allocator;
3const sort = std.sort;
4const assert = std.debug.assert;
5const warn = std.debug.warn;
6const testing = std.testing;
7const expect = testing.expect;
8const expectEqual = testing.expectEqual;
9const expectError = testing.expectError;
10
11/// Priority Dequeue for storing generic data. Initialize with `init`.
12pub fn PriorityDequeue(comptime T: type) type {
13 return struct {
14 const Self = @This();
15
16 items: []T,
17 len: usize,
18 allocator: *Allocator,
19 lessThanFn: fn (a: T, b: T) bool,
20
21 /// Initialize and return a new dequeue. Provide `lessThanFn`
22 /// that returns `true` when its first argument should
23 /// get min-popped before its second argument. For example,
24 /// to make `popMin` return the minimum value, provide
25 ///
26 /// `fn lessThanFn(a: T, b: T) bool { return a < b; }`
27 pub fn init(allocator: *Allocator, lessThanFn: fn (T, T) bool) Self {
28 return Self{
29 .items = &[_]T{},
30 .len = 0,
31 .allocator = allocator,
32 .lessThanFn = lessThanFn,
33 };
34 }
35
36 fn lessThan(self: Self, a: T, b: T) bool {
37 return self.lessThanFn(a, b);
38 }
39
40 fn greaterThan(self: Self, a: T, b: T) bool {
41 return self.lessThanFn(b, a);
42 }
43
44 /// Free memory used by the dequeue.
45 pub fn deinit(self: Self) void {
46 self.allocator.free(self.items);
47 }
48
49 /// Insert a new element, maintaining priority.
50 pub fn add(self: *Self, elem: T) !void {
51 try ensureCapacity(self, self.len + 1);
52 addUnchecked(self, elem);
53 }
54
55 /// Add each element in `items` to the dequeue.
56 pub fn addSlice(self: *Self, items: []const T) !void {
57 try self.ensureCapacity(self.len + items.len);
58 for (items) |e| {
59 self.addUnchecked(e);
60 }
61 }
62
63 fn addUnchecked(self: *Self, elem: T) void {
64 self.items[self.len] = elem;
65
66 if (self.len > 0) {
67 const start = self.getStartForSiftUp(elem, self.len);
68 self.siftUp(start);
69 }
70
71 self.len += 1;
72 }
73
74 fn isMinLayer(index: usize) bool {
75 // In the min-max heap structure:
76 // The first element is on a min layer;
77 // next two are on a max layer;
78 // next four are on a min layer, and so on.
79 const leading_zeros = @clz(usize, index + 1);
80 const highest_set_bit = 63 - leading_zeros;
81 return (highest_set_bit & 1) == 0;
82 }
83
84 fn nextIsMinLayer(self: Self) bool {
85 return isMinLayer(self.len);
86 }
87
88 const StartIndexAndLayer = struct {
89 index: usize,
90 min_layer: bool,
91 };
92
93 fn getStartForSiftUp(self: Self, child: T, index: usize) StartIndexAndLayer {
94 var child_index = index;
95 var parent_index = parentIndex(child_index);
96 const parent = self.items[parent_index];
97
98 const min_layer = self.nextIsMinLayer();
99 if ((min_layer and self.greaterThan(child, parent)) or (!min_layer and self.lessThan(child, parent))) {
100 // We must swap the item with it's parent if it is on the "wrong" layer
101 self.items[parent_index] = child;
102 self.items[child_index] = parent;
103 return .{
104 .index = parent_index,
105 .min_layer = !min_layer,
106 };
107 } else {
108 return .{
109 .index = child_index,
110 .min_layer = min_layer,
111 };
112 }
113 }
114
115 fn siftUp(self: *Self, start: StartIndexAndLayer) void {
116 if (start.min_layer) {
117 doSiftUp(self, start.index, lessThan);
118 } else {
119 doSiftUp(self, start.index, greaterThan);
120 }
121 }
122
123 fn doSiftUp(self: *Self, start_index: usize, compare: fn (Self, T, T) bool) void {
124 var child_index = start_index;
125 while (child_index > 2) {
126 var grandparent_index = grandparentIndex(child_index);
127 const child = self.items[child_index];
128 const grandparent = self.items[grandparent_index];
129
130 // If the grandparent is already better, we have gone as far as we need to
131 if (!compare(self.*, child, grandparent)) break;
132
133 // Otherwise swap the item with it's grandparent
134 self.items[grandparent_index] = child;
135 self.items[child_index] = grandparent;
136 child_index = grandparent_index;
137 }
138 }
139
140 /// Look at the smallest element in the dequeue. Returns
141 /// `null` if empty.
142 pub fn peekMin(self: *Self) ?T {
143 return if (self.len > 0) self.items[0] else null;
144 }
145
146 /// Look at the largest element in the dequeue. Returns
147 /// `null` if empty.
148 pub fn peekMax(self: *Self) ?T {
149 if (self.len == 0) return null;
150 if (self.len == 1) return self.items[0];
151 if (self.len == 2) return self.items[1];
152 return self.bestItemAtIndices(1, 2, greaterThan).item;
153 }
154
155 fn maxIndex(self: Self) ?usize {
156 if (self.len == 0) return null;
157 if (self.len == 1) return 0;
158 if (self.len == 2) return 1;
159 return self.bestItemAtIndices(1, 2, greaterThan).index;
160 }
161
162 /// Pop the smallest element from the dequeue. Returns
163 /// `null` if empty.
164 pub fn removeMinOrNull(self: *Self) ?T {
165 return if (self.len > 0) self.removeMin() else null;
166 }
167
168 /// Remove and return the smallest element from the
169 /// dequeue.
170 pub fn removeMin(self: *Self) T {
171 return self.removeIndex(0);
172 }
173
174 /// Pop the largest element from the dequeue. Returns
175 /// `null` if empty.
176 pub fn removeMaxOrNull(self: *Self) ?T {
177 return if (self.len > 0) self.removeMax() else null;
178 }
179
180 /// Remove and return the largest element from the
181 /// dequeue.
182 pub fn removeMax(self: *Self) T {
183 return self.removeIndex(self.maxIndex().?);
184 }
185
186 /// Remove and return element at index. Indices are in the
187 /// same order as iterator, which is not necessarily priority
188 /// order.
189 pub fn removeIndex(self: *Self, index: usize) T {
190 const item = self.items[index];
191 const last = self.items[self.len - 1];
192
193 self.items[index] = last;
194 self.len -= 1;
195 siftDown(self, index);
196
197 return item;
198 }
199
200 fn siftDown(self: *Self, index: usize) void {
201 if (isMinLayer(index)) {
202 self.doSiftDown(index, lessThan);
203 } else {
204 self.doSiftDown(index, greaterThan);
205 }
206 }
207
208 fn doSiftDown(self: *Self, start_index: usize, compare: fn (Self, T, T) bool) void {
209 var index = start_index;
210 const half = self.len >> 1;
211 while (true) {
212 const first_grandchild_index = firstGrandchildIndex(index);
213 const last_grandchild_index = first_grandchild_index + 3;
214
215 const elem = self.items[index];
216
217 if (last_grandchild_index < self.len) {
218 // All four grandchildren exist
219 const index2 = first_grandchild_index + 1;
220 const index3 = index2 + 1;
221
222 // Find the best grandchild
223 const best_left = self.bestItemAtIndices(first_grandchild_index, index2, compare);
224 const best_right = self.bestItemAtIndices(index3, last_grandchild_index, compare);
225 const best_grandchild = self.bestItem(best_left, best_right, compare);
226
227 // If the item is better than it's best grandchild, we are done
228 if (compare(self.*, elem, best_grandchild.item) or elem == best_grandchild.item) return;
229
230 // Otherwise, swap them
231 self.items[best_grandchild.index] = elem;
232 self.items[index] = best_grandchild.item;
233 index = best_grandchild.index;
234
235 // We might need to swap the element with it's parent
236 self.swapIfParentIsBetter(elem, index, compare);
237 } else {
238 // The children or grandchildren are the last layer
239 const first_child_index = firstChildIndex(index);
240 if (first_child_index > self.len) return;
241
242 const best_descendent = self.bestDescendent(first_child_index, first_grandchild_index, compare);
243
244 // If the best descendant is still larger, we are done
245 if (compare(self.*, elem, best_descendent.item) or elem == best_descendent.item) return;
246
247 // Otherwise swap them
248 self.items[best_descendent.index] = elem;
249 self.items[index] = best_descendent.item;
250 index = best_descendent.index;
251
252 // If we didn't swap a grandchild, we are done
253 if (index < first_grandchild_index) return;
254
255 // We might need to swap the element with it's parent
256 self.swapIfParentIsBetter(elem, index, compare);
257 return;
258 }
259
260 // If we are now in the last layer, we are done
261 if (index >= half) return;
262 }
263 }
264
265 fn swapIfParentIsBetter(self: *Self, child: T, child_index: usize, compare: fn (Self, T, T) bool) void {
266 const parent_index = parentIndex(child_index);
267 const parent = self.items[parent_index];
268
269 if (compare(self.*, parent, child)) {
270 self.items[parent_index] = child;
271 self.items[child_index] = parent;
272 }
273 }
274
275 const ItemAndIndex = struct {
276 item: T,
277 index: usize,
278 };
279
280 fn getItem(self: Self, index: usize) ItemAndIndex {
281 return .{
282 .item = self.items[index],
283 .index = index,
284 };
285 }
286
287 fn bestItem(self: Self, item1: ItemAndIndex, item2: ItemAndIndex, compare: fn (Self, T, T) bool) ItemAndIndex {
288 if (compare(self, item1.item, item2.item)) {
289 return item1;
290 } else {
291 return item2;
292 }
293 }
294
295 fn bestItemAtIndices(self: Self, index1: usize, index2: usize, compare: fn (Self, T, T) bool) ItemAndIndex {
296 var item1 = self.getItem(index1);
297 var item2 = self.getItem(index2);
298 return self.bestItem(item1, item2, compare);
299 }
300
301 fn bestDescendent(self: Self, first_child_index: usize, first_grandchild_index: usize, compare: fn (Self, T, T) bool) ItemAndIndex {
302 const second_child_index = first_child_index + 1;
303 if (first_grandchild_index >= self.len) {
304 // No grandchildren, find the best child (second may not exist)
305 if (second_child_index >= self.len) {
306 return .{
307 .item = self.items[first_child_index],
308 .index = first_child_index,
309 };
310 } else {
311 return self.bestItemAtIndices(first_child_index, second_child_index, compare);
312 }
313 }
314
315 const second_grandchild_index = first_grandchild_index + 1;
316 if (second_grandchild_index >= self.len) {
317 // One grandchild, so we know there is a second child. Compare first grandchild and second child
318 return self.bestItemAtIndices(first_grandchild_index, second_child_index, compare);
319 }
320
321 const best_left_grandchild_index = self.bestItemAtIndices(first_grandchild_index, second_grandchild_index, compare).index;
322 const third_grandchild_index = second_grandchild_index + 1;
323 if (third_grandchild_index >= self.len) {
324 // Two grandchildren, and we know the best. Compare this to second child.
325 return self.bestItemAtIndices(best_left_grandchild_index, second_child_index, compare);
326 } else {
327 // Three grandchildren, compare the min of the first two with the third
328 return self.bestItemAtIndices(best_left_grandchild_index, third_grandchild_index, compare);
329 }
330 }
331
332 /// Return the number of elements remaining in the heap
333 pub fn count(self: Self) usize {
334 return self.len;
335 }
336
337 /// Return the number of elements that can be added to the
338 /// dequeue before more memory is allocated.
339 pub fn capacity(self: Self) usize {
340 return self.items.len;
341 }
342
343 /// Heap takes ownership of the passed in slice. The slice must have been
344 /// allocated with `allocator`.
345 /// De-initialize with `deinit`.
346 pub fn fromOwnedSlice(allocator: *Allocator, lessThanFn: fn (T, T) bool, items: []T) Self {
347 var dequeue = Self{
348 .items = items,
349 .len = items.len,
350 .allocator = allocator,
351 .lessThanFn = lessThanFn,
352 };
353 const half = (dequeue.len >> 1) - 1;
354 var i: usize = 0;
355 while (i <= half) : (i += 1) {
356 const index = half - i;
357 dequeue.siftDown(index);
358 }
359 return dequeue;
360 }
361
362 pub fn ensureCapacity(self: *Self, new_capacity: usize) !void {
363 var better_capacity = self.capacity();
364 if (better_capacity >= new_capacity) return;
365 while (true) {
366 better_capacity += better_capacity / 2 + 8;
367 if (better_capacity >= new_capacity) break;
368 }
369 self.items = try self.allocator.realloc(self.items, better_capacity);
370 }
371
372 pub fn resize(self: *Self, new_len: usize) !void {
373 try self.ensureCapacity(new_len);
374 self.len = new_len;
375 }
376
377 pub fn shrink(self: *Self, new_len: usize) void {
378 // TODO take advantage of the new realloc semantics
379 assert(new_len <= self.len);
380 self.len = new_len;
381 }
382
383 pub fn update(self: *Self, elem: T, new_elem: T) !void {
384 var old_index: usize = std.mem.indexOfScalar(T, self.items, elem) orelse return error.ElementNotFound;
385 _ = self.removeIndex(old_index);
386 self.addUnchecked(new_elem);
387 }
388
389 pub const Iterator = struct {
390 heap: *PriorityDequeue(T),
391 count: usize,
392
393 pub fn next(it: *Iterator) ?T {
394 if (it.count >= it.heap.len) return null;
395 const out = it.count;
396 it.count += 1;
397 return it.heap.items[out];
398 }
399
400 pub fn reset(it: *Iterator) void {
401 it.count = 0;
402 }
403 };
404
405 /// Return an iterator that walks the heap without consuming
406 /// it. Invalidated if the heap is modified.
407 pub fn iterator(self: *Self) Iterator {
408 return Iterator{
409 .heap = self,
410 .count = 0,
411 };
412 }
413
414 fn dump(self: *Self) void {
415 warn("{{ ", .{});
416 warn("items: ", .{});
417 for (self.items) |e, i| {
418 if (i >= self.len) break;
419 warn("{}, ", .{e});
420 }
421 warn("array: ", .{});
422 for (self.items) |e, i| {
423 warn("{}, ", .{e});
424 }
425 warn("len: {} ", .{self.len});
426 warn("capacity: {}", .{self.capacity()});
427 warn(" }}\n", .{});
428 }
429
430 fn parentIndex(index: usize) usize {
431 return (index - 1) >> 1;
432 }
433
434 fn grandparentIndex(index: usize) usize {
435 return parentIndex(parentIndex(index));
436 }
437
438 fn firstChildIndex(index: usize) usize {
439 return (index << 1) + 1;
440 }
441
442 fn firstGrandchildIndex(index: usize) usize {
443 return firstChildIndex(firstChildIndex(index));
444 }
445 };
446}
447
448fn lessThanComparison(a: u32, b: u32) bool {
449 return a < b;
450}
451
452const Heap = PriorityDequeue(u32);
453
454test "std.PriorityDequeue: add and remove min" {
455 var heap = Heap.init(testing.allocator, lessThanComparison);
456 defer heap.deinit();
457
458 try heap.add(54);
459 try heap.add(12);
460 try heap.add(7);
461 try heap.add(23);
462 try heap.add(25);
463 try heap.add(13);
464
465 expectEqual(@as(u32, 7), heap.removeMin());
466 expectEqual(@as(u32, 12), heap.removeMin());
467 expectEqual(@as(u32, 13), heap.removeMin());
468 expectEqual(@as(u32, 23), heap.removeMin());
469 expectEqual(@as(u32, 25), heap.removeMin());
470 expectEqual(@as(u32, 54), heap.removeMin());
471}
472
473test "std.PriorityDequeue: add and remove max" {
474 var heap = Heap.init(testing.allocator, lessThanComparison);
475 defer heap.deinit();
476
477 try heap.add(54);
478 try heap.add(12);
479 try heap.add(7);
480 try heap.add(23);
481 try heap.add(25);
482 try heap.add(13);
483
484 expectEqual(@as(u32, 54), heap.removeMax());
485 expectEqual(@as(u32, 25), heap.removeMax());
486 expectEqual(@as(u32, 23), heap.removeMax());
487 expectEqual(@as(u32, 13), heap.removeMax());
488 expectEqual(@as(u32, 12), heap.removeMax());
489 expectEqual(@as(u32, 7), heap.removeMax());
490}
491
492test "std.PriorityDequeue: add and remove same min" {
493 var heap = Heap.init(testing.allocator, lessThanComparison);
494 defer heap.deinit();
495
496 try heap.add(1);
497 try heap.add(1);
498 try heap.add(2);
499 try heap.add(2);
500 try heap.add(1);
501 try heap.add(1);
502
503 expectEqual(@as(u32, 1), heap.removeMin());
504 expectEqual(@as(u32, 1), heap.removeMin());
505 expectEqual(@as(u32, 1), heap.removeMin());
506 expectEqual(@as(u32, 1), heap.removeMin());
507 expectEqual(@as(u32, 2), heap.removeMin());
508 expectEqual(@as(u32, 2), heap.removeMin());
509}
510
511test "std.PriorityDequeue: add and remove same max" {
512 var heap = Heap.init(testing.allocator, lessThanComparison);
513 defer heap.deinit();
514
515 try heap.add(1);
516 try heap.add(1);
517 try heap.add(2);
518 try heap.add(2);
519 try heap.add(1);
520 try heap.add(1);
521
522 expectEqual(@as(u32, 2), heap.removeMax());
523 expectEqual(@as(u32, 2), heap.removeMax());
524 expectEqual(@as(u32, 1), heap.removeMax());
525 expectEqual(@as(u32, 1), heap.removeMax());
526 expectEqual(@as(u32, 1), heap.removeMax());
527 expectEqual(@as(u32, 1), heap.removeMax());
528}
529
530test "std.PriorityDequeue: removeOrNull empty" {
531 var heap = Heap.init(testing.allocator, lessThanComparison);
532 defer heap.deinit();
533
534 expect(heap.removeMinOrNull() == null);
535 expect(heap.removeMaxOrNull() == null);
536}
537
538test "std.PriorityDequeue: edge case 3 elements" {
539 var heap = Heap.init(testing.allocator, lessThanComparison);
540 defer heap.deinit();
541
542 try heap.add(9);
543 try heap.add(3);
544 try heap.add(2);
545
546 expectEqual(@as(u32, 2), heap.removeMin());
547 expectEqual(@as(u32, 3), heap.removeMin());
548 expectEqual(@as(u32, 9), heap.removeMin());
549}
550
551test "std.PriorityDequeue: edge case 3 elements max" {
552 var heap = Heap.init(testing.allocator, lessThanComparison);
553 defer heap.deinit();
554
555 try heap.add(9);
556 try heap.add(3);
557 try heap.add(2);
558
559 expectEqual(@as(u32, 9), heap.removeMax());
560 expectEqual(@as(u32, 3), heap.removeMax());
561 expectEqual(@as(u32, 2), heap.removeMax());
562}
563
564test "std.PriorityDequeue: peekMin" {
565 var heap = Heap.init(testing.allocator, lessThanComparison);
566 defer heap.deinit();
567
568 expect(heap.peekMin() == null);
569
570 try heap.add(9);
571 try heap.add(3);
572 try heap.add(2);
573
574 expect(heap.peekMin().? == 2);
575 expect(heap.peekMin().? == 2);
576}
577
578test "std.PriorityDequeue: peekMax" {
579 var heap = Heap.init(testing.allocator, lessThanComparison);
580 defer heap.deinit();
581
582 expect(heap.peekMin() == null);
583
584 try heap.add(9);
585 try heap.add(3);
586 try heap.add(2);
587
588 expect(heap.peekMax().? == 9);
589 expect(heap.peekMax().? == 9);
590}
591
592test "std.PriorityDequeue: sift up with odd indices" {
593 var heap = Heap.init(testing.allocator, lessThanComparison);
594 defer heap.deinit();
595 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
596 for (items) |e| {
597 try heap.add(e);
598 }
599
600 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
601 for (sorted_items) |e| {
602 expectEqual(e, heap.removeMin());
603 }
604}
605
606test "std.PriorityDequeue: sift up with odd indices" {
607 var heap = Heap.init(testing.allocator, lessThanComparison);
608 defer heap.deinit();
609 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
610 for (items) |e| {
611 try heap.add(e);
612 }
613
614 const sorted_items = [_]u32{ 25, 24, 24, 22, 21, 16, 15, 15, 14, 13, 12, 11, 7, 7, 6, 5, 2, 1 };
615 for (sorted_items) |e| {
616 expectEqual(e, heap.removeMax());
617 }
618}
619
620test "std.PriorityDequeue: addSlice min" {
621 var heap = Heap.init(testing.allocator, lessThanComparison);
622 defer heap.deinit();
623 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
624 try heap.addSlice(items[0..]);
625
626 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
627 for (sorted_items) |e| {
628 expectEqual(e, heap.removeMin());
629 }
630}
631
632test "std.PriorityDequeue: addSlice max" {
633 var heap = Heap.init(testing.allocator, lessThanComparison);
634 defer heap.deinit();
635 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
636 try heap.addSlice(items[0..]);
637
638 const sorted_items = [_]u32{ 25, 24, 24, 22, 21, 16, 15, 15, 14, 13, 12, 11, 7, 7, 6, 5, 2, 1 };
639 for (sorted_items) |e| {
640 expectEqual(e, heap.removeMax());
641 }
642}
643
644test "std.PriorityDequeue: fromOwnedSlice" {
645 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
646 const heap_items = try testing.allocator.dupe(u32, items[0..]);
647 var heap = Heap.fromOwnedSlice(testing.allocator, lessThanComparison, heap_items[0..]);
648 defer heap.deinit();
649
650 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
651 for (sorted_items) |e| {
652 expectEqual(e, heap.removeMin());
653 }
654}
655
656test "std.PriorityDequeue: update min heap" {
657 var heap = Heap.init(testing.allocator, lessThanComparison);
658 defer heap.deinit();
659
660 try heap.add(55);
661 try heap.add(44);
662 try heap.add(11);
663 try heap.update(55, 5);
664 try heap.update(44, 4);
665 try heap.update(11, 1);
666 expectEqual(@as(u32, 1), heap.removeMin());
667 expectEqual(@as(u32, 4), heap.removeMin());
668 expectEqual(@as(u32, 5), heap.removeMin());
669}
670
671test "std.PriorityDequeue: update same min heap" {
672 var heap = Heap.init(testing.allocator, lessThanComparison);
673 defer heap.deinit();
674
675 try heap.add(1);
676 try heap.add(1);
677 try heap.add(2);
678 try heap.add(2);
679 try heap.update(1, 5);
680 try heap.update(2, 4);
681 expectEqual(@as(u32, 1), heap.removeMin());
682 expectEqual(@as(u32, 2), heap.removeMin());
683 expectEqual(@as(u32, 4), heap.removeMin());
684 expectEqual(@as(u32, 5), heap.removeMin());
685}
686
687test "std.PriorityDequeue: update max heap" {
688 var heap = Heap.init(testing.allocator, lessThanComparison);
689 defer heap.deinit();
690
691 try heap.add(55);
692 try heap.add(44);
693 try heap.add(11);
694 try heap.update(55, 5);
695 try heap.update(44, 1);
696 try heap.update(11, 4);
697
698 expectEqual(@as(u32, 5), heap.removeMax());
699 expectEqual(@as(u32, 4), heap.removeMax());
700 expectEqual(@as(u32, 1), heap.removeMax());
701}
702
703test "std.PriorityDequeue: update same max heap" {
704 var heap = Heap.init(testing.allocator, lessThanComparison);
705 defer heap.deinit();
706
707 try heap.add(1);
708 try heap.add(1);
709 try heap.add(2);
710 try heap.add(2);
711 try heap.update(1, 5);
712 try heap.update(2, 4);
713 expectEqual(@as(u32, 5), heap.removeMax());
714 expectEqual(@as(u32, 4), heap.removeMax());
715 expectEqual(@as(u32, 2), heap.removeMax());
716 expectEqual(@as(u32, 1), heap.removeMax());
717}
718
719test "std.PriorityDequeue: iterator" {
720 var heap = Heap.init(testing.allocator, lessThanComparison);
721 var map = std.AutoHashMap(u32, void).init(testing.allocator);
722 defer {
723 heap.deinit();
724 map.deinit();
725 }
726
727 const items = [_]u32{ 54, 12, 7, 23, 25, 13 };
728 for (items) |e| {
729 _ = try heap.add(e);
730 _ = try map.put(e, {});
731 }
732
733 var it = heap.iterator();
734 while (it.next()) |e| {
735 _ = map.remove(e);
736 }
737
738 expectEqual(@as(usize, 0), map.count());
739}
740
741test "std.PriorityDequeue: remove at index" {
742 var heap = Heap.init(testing.allocator, lessThanComparison);
743 defer heap.deinit();
744
745 try heap.add(3);
746 try heap.add(2);
747 try heap.add(1);
748
749 var it = heap.iterator();
750 var elem = it.next();
751 var idx: usize = 0;
752 const two_idx = while (elem != null) : (elem = it.next()) {
753 if (elem.? == 2)
754 break idx;
755 idx += 1;
756 } else unreachable;
757
758 expectEqual(heap.removeIndex(two_idx), 2);
759 expectEqual(heap.removeMin(), 1);
760 expectEqual(heap.removeMin(), 3);
761 expectEqual(heap.removeMinOrNull(), null);
762}
763
764test "std.PriorityDequeue: iterator while empty" {
765 var heap = Heap.init(testing.allocator, lessThanComparison);
766 defer heap.deinit();
767
768 var it = heap.iterator();
769
770 expectEqual(it.next(), null);
771}
772
773test "std.PriorityDequeue: fuzz testing min" {
774 var prng = std.rand.DefaultPrng.init(0x12345678);
775
776 const test_case_count = 100;
777 const heap_size = 1_000;
778
779 var i: usize = 0;
780 while (i < test_case_count) : (i += 1) {
781 try fuzzTestMin(&prng.random, heap_size);
782 }
783}
784
785fn fuzzTestMin(rng: *std.rand.Random, comptime heap_size: usize) !void {
786 const allocator = testing.allocator;
787 const items = try generateRandomSlice(allocator, rng, heap_size);
788
789 var heap = Heap.fromOwnedSlice(allocator, lessThanComparison, items);
790 defer heap.deinit();
791
792 var last_removed: ?u32 = null;
793 while (heap.removeMinOrNull()) |next| {
794 if (last_removed) |last| {
795 expect(last <= next);
796 }
797 last_removed = next;
798 }
799}
800
801test "std.PriorityDequeue: fuzz testing max" {
802 var prng = std.rand.DefaultPrng.init(0x87654321);
803
804 const test_case_count = 100;
805 const heap_size = 1_000;
806
807 var i: usize = 0;
808 while (i < test_case_count) : (i += 1) {
809 try fuzzTestMax(&prng.random, heap_size);
810 }
811}
812
813fn fuzzTestMax(rng: *std.rand.Random, heap_size: usize) !void {
814 const allocator = testing.allocator;
815 const items = try generateRandomSlice(allocator, rng, heap_size);
816
817 var heap = Heap.fromOwnedSlice(testing.allocator, lessThanComparison, items);
818 defer heap.deinit();
819
820 var last_removed: ?u32 = null;
821 while (heap.removeMaxOrNull()) |next| {
822 if (last_removed) |last| {
823 expect(last >= next);
824 }
825 last_removed = next;
826 }
827}
828
829test "std.PriorityDequeue: fuzz testing min and max" {
830 var prng = std.rand.DefaultPrng.init(0x87654321);
831
832 const test_case_count = 100;
833 const heap_size = 1_000;
834
835 var i: usize = 0;
836 while (i < test_case_count) : (i += 1) {
837 try fuzzTestMinMax(&prng.random, heap_size);
838 }
839}
840
841fn fuzzTestMinMax(rng: *std.rand.Random, heap_size: usize) !void {
842 const allocator = testing.allocator;
843 const items = try generateRandomSlice(allocator, rng, heap_size);
844
845 var heap = Heap.fromOwnedSlice(allocator, lessThanComparison, items);
846 defer heap.deinit();
847
848 var last_min: ?u32 = null;
849 var last_max: ?u32 = null;
850 var i: usize = 0;
851 while (i < heap_size) : (i += 1) {
852 if (i % 2 == 0) {
853 const next = heap.removeMin();
854 if (last_min) |last| {
855 expect(last <= next);
856 }
857 last_min = next;
858 } else {
859 const next = heap.removeMax();
860 if (last_max) |last| {
861 expect(last >= next);
862 }
863 last_max = next;
864 }
865 }
866}
867
868fn generateRandomSlice(allocator: *std.mem.Allocator, rng: *std.rand.Random, size: usize) ![]u32 {
869 var array = std.ArrayList(u32).init(allocator);
870 try array.ensureCapacity(size);
871
872 var i: usize = 0;
873 while (i < size) : (i += 1) {
874 const elem = rng.int(u32);
875 try array.append(elem);
876 }
877
878 return array.toOwnedSlice();
879}
lib/std/std.zig+1
......@@ -25,6 +25,7 @@ pub const PackedIntArrayEndian = @import("packed_int_array.zig").PackedIntArrayE
2525pub const PackedIntSlice = @import("packed_int_array.zig").PackedIntSlice;
2626pub const PackedIntSliceEndian = @import("packed_int_array.zig").PackedIntSliceEndian;
2727pub const PriorityQueue = @import("priority_queue.zig").PriorityQueue;
28pub const PriorityDequeue = @import("priority_dequeue.zig").PriorityDequeue;
2829pub const Progress = @import("Progress.zig");
2930pub const SemanticVersion = @import("SemanticVersion.zig");
3031pub const SinglyLinkedList = @import("linked_list.zig").SinglyLinkedList;