1const std = @import("std.zig");
2const Allocator = std.mem.Allocator;
3const assert = std.debug.assert;
4const Order = std.math.Order;
5const testing = std.testing;
6const expect = testing.expect;
7const expectEqual = testing.expectEqual;
8const expectError = testing.expectError;
9
10/// Priority queue for storing generic data. Initialize with `init`.
11/// Provide `compareFn` that returns `Order.lt` when its second
12/// argument should get popped before its third argument,
13/// `Order.eq` if the arguments are of equal priority, or `Order.gt`
14/// if the third argument should be popped first.
15/// For example, to make `pop` return the smallest number, provide
16/// `fn lessThan(context: void, a: T, b: T) Order { _ = context; return std.math.order(a, b); }`
17pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareFn: fn (context: Context, a: T, b: T) Order) type {
18 return struct {
19 const Self = @This();
20
21 items: []T,
22 cap: usize,
23 context: Context,
24
25 /// A priority queue containing no elements.
26 pub const empty: Self = .{
27 .items = &.{},
28 .cap = 0,
29 .context = undefined,
30 };
31
32 /// Initialize and return a priority queue with context.
33 pub fn initContext(context: Context) Self {
34 return Self{
35 .items = &.{},
36 .cap = 0,
37 .context = context,
38 };
39 }
40
41 /// Free memory used by the queue.
42 pub fn deinit(self: *Self, allocator: Allocator) void {
43 allocator.free(self.allocatedSlice());
44 self.* = undefined;
45 }
46
47 /// Insert a new element, maintaining priority.
48 pub fn push(self: *Self, allocator: Allocator, elem: T) !void {
49 try self.ensureUnusedCapacity(allocator, 1);
50 pushUnchecked(self, elem);
51 }
52
53 fn pushUnchecked(self: *Self, elem: T) void {
54 self.items.len += 1;
55 self.items[self.items.len - 1] = elem;
56 siftUp(self, self.items.len - 1);
57 }
58
59 fn siftUp(self: *Self, start_index: usize) void {
60 const child = self.items[start_index];
61 var child_index = start_index;
62 while (child_index > 0) {
63 const parent_index = ((child_index - 1) >> 1);
64 const parent = self.items[parent_index];
65 if (compareFn(self.context, child, parent) != .lt) break;
66 self.items[child_index] = parent;
67 child_index = parent_index;
68 }
69 self.items[child_index] = child;
70 }
71
72 /// Add each element in `items` to the queue.
73 pub fn pushSlice(self: *Self, allocator: Allocator, items: []const T) !void {
74 try self.ensureUnusedCapacity(allocator, items.len);
75 for (items) |e| {
76 self.pushUnchecked(e);
77 }
78 }
79
80 /// Look at the highest priority element in the queue. Returns
81 /// `null` if empty.
82 pub fn peek(self: *const Self) ?T {
83 return if (self.items.len > 0) self.items[0] else null;
84 }
85
86 /// Remove and return the highest priority element from the queue.
87 /// Returns `null` if empty.
88 pub fn pop(self: *Self) ?T {
89 return if (self.items.len > 0) self.popIndex(0) else null;
90 }
91
92 /// Remove and return element at index. Indices are in the
93 /// same order as iterator, which is not necessarily priority
94 /// order.
95 pub fn popIndex(self: *Self, index: usize) T {
96 assert(self.items.len > index);
97 const last = self.items[self.items.len - 1];
98 const item = self.items[index];
99 self.items[index] = last;
100 self.items.len -= 1;
101
102 if (index == self.items.len) {
103 // Last element removed, nothing more to do.
104 } else if (index == 0) {
105 siftDown(self, index);
106 } else {
107 const parent_index = ((index - 1) >> 1);
108 const parent = self.items[parent_index];
109 if (compareFn(self.context, last, parent) == .gt) {
110 siftDown(self, index);
111 } else {
112 siftUp(self, index);
113 }
114 }
115
116 return item;
117 }
118
119 /// Return the number of elements remaining in the priority
120 /// queue.
121 pub fn count(self: *const Self) usize {
122 return self.items.len;
123 }
124
125 /// Return the number of elements that can be added to the
126 /// queue before more memory is allocated.
127 pub fn capacity(self: *const Self) usize {
128 return self.cap;
129 }
130
131 /// Returns a slice of all the items plus the extra capacity, whose memory
132 /// contents are `undefined`.
133 fn allocatedSlice(self: *const Self) []T {
134 // `items.len` is the length, not the capacity.
135 return self.items.ptr[0..self.cap];
136 }
137
138 fn siftDown(self: *Self, target_index: usize) void {
139 const target_element = self.items[target_index];
140 var index = target_index;
141 while (true) {
142 var lesser_child_i = (std.math.mul(usize, index, 2) catch break) | 1;
143 if (!(lesser_child_i < self.items.len)) break;
144
145 const next_child_i = lesser_child_i + 1;
146 if (next_child_i < self.items.len and compareFn(self.context, self.items[next_child_i], self.items[lesser_child_i]) == .lt) {
147 lesser_child_i = next_child_i;
148 }
149
150 if (compareFn(self.context, target_element, self.items[lesser_child_i]) == .lt) break;
151
152 self.items[index] = self.items[lesser_child_i];
153 index = lesser_child_i;
154 }
155 self.items[index] = target_element;
156 }
157
158 /// PriorityQueue takes ownership of the passed in slice. The slice must have been
159 /// allocated with `allocator`.
160 /// Deinitialize with `deinit`.
161 pub fn fromOwnedSlice(items: []T, context: Context) Self {
162 var self = Self{
163 .items = items,
164 .cap = items.len,
165 .context = context,
166 };
167
168 var i = self.items.len >> 1;
169 while (i > 0) {
170 i -= 1;
171 self.siftDown(i);
172 }
173 return self;
174 }
175
176 /// Ensure that the queue can fit at least `new_capacity` items.
177 pub fn ensureTotalCapacity(self: *Self, allocator: Allocator, new_capacity: usize) !void {
178 var better_capacity = self.cap;
179 if (better_capacity >= new_capacity) return;
180 while (true) {
181 better_capacity += better_capacity / 2 + 8;
182 if (better_capacity >= new_capacity) break;
183 }
184 try self.ensureTotalCapacityPrecise(allocator, better_capacity);
185 }
186
187 /// If the current capacity is less than `new_capacity`, this function will
188 /// modify the array so that it can hold exactly `new_capacity` items.
189 /// Invalidates element pointers if additional memory is needed.
190 pub fn ensureTotalCapacityPrecise(self: *Self, allocator: Allocator, new_capacity: usize) !void {
191 if (self.capacity() >= new_capacity) return;
192
193 const old_memory = self.allocatedSlice();
194 const new_memory = try allocator.realloc(old_memory, new_capacity);
195 self.items.ptr = new_memory.ptr;
196 self.cap = new_memory.len;
197 }
198
199 /// Ensure that the queue can fit at least `additional_count` **more** item.
200 pub fn ensureUnusedCapacity(self: *Self, allocator: Allocator, additional_count: usize) !void {
201 return self.ensureTotalCapacity(allocator, self.items.len + additional_count);
202 }
203
204 /// Reduce allocated capacity to `new_capacity`.
205 pub fn shrinkAndFree(self: *Self, allocator: Allocator, new_capacity: usize) void {
206 assert(new_capacity <= self.cap);
207
208 // Cannot shrink to smaller than the current queue size without invalidating the heap property
209 assert(new_capacity >= self.items.len);
210
211 const old_memory = self.allocatedSlice();
212 const new_memory = allocator.realloc(old_memory, new_capacity) catch |e| switch (e) {
213 error.OutOfMemory => { // no problem, capacity is still correct then.
214 return;
215 },
216 };
217
218 self.items.ptr = new_memory.ptr;
219 self.cap = new_memory.len;
220 }
221
222 /// Remove all elements from the items slice.
223 pub fn clearRetainingCapacity(self: *Self) void {
224 self.items.len = 0;
225 }
226
227 /// Invalidates all element pointers.
228 pub fn clearAndFree(self: *Self, allocator: Allocator) void {
229 allocator.free(self.allocatedSlice());
230 self.items.len = 0;
231 self.cap = 0;
232 }
233
234 /// Replace an element in the queue with a new element, maintaining priority.
235 /// If the element being updated doesn't exist, return `error.ElementNotFound`.
236 pub fn update(self: *Self, elem: T, new_elem: T) !void {
237 const update_index = blk: {
238 var idx: usize = 0;
239 while (idx < self.items.len) : (idx += 1) {
240 const item = self.items[idx];
241 if (compareFn(self.context, item, elem) == .eq) break :blk idx;
242 }
243 return error.ElementNotFound;
244 };
245 const old_elem: T = self.items[update_index];
246 self.items[update_index] = new_elem;
247 switch (compareFn(self.context, new_elem, old_elem)) {
248 .lt => siftUp(self, update_index),
249 .gt => siftDown(self, update_index),
250 .eq => {}, // Nothing to do as the items have equal priority
251 }
252 }
253
254 pub const Iterator = struct {
255 queue: *PriorityQueue(T, Context, compareFn),
256 count: usize,
257
258 pub fn next(it: *Iterator) ?T {
259 if (it.count >= it.queue.items.len) return null;
260 const out = it.count;
261 it.count += 1;
262 return it.queue.items[out];
263 }
264
265 pub fn reset(it: *Iterator) void {
266 it.count = 0;
267 }
268 };
269
270 /// Return an iterator that walks the queue without consuming
271 /// it. The iteration order may differ from the priority order.
272 /// Invalidated if the heap is modified.
273 pub fn iterator(self: *Self) Iterator {
274 return Iterator{
275 .queue = self,
276 .count = 0,
277 };
278 }
279 };
280}
281
282fn lessThan(context: void, a: u32, b: u32) Order {
283 _ = context;
284 return std.math.order(a, b);
285}
286
287fn greaterThan(context: void, a: u32, b: u32) Order {
288 return lessThan(context, a, b).invert();
289}
290
291const MinHeap = PriorityQueue(u32, void, lessThan);
292const MaxHeap = PriorityQueue(u32, void, greaterThan);
293
294test "add and remove min heap" {
295 const gpa = testing.allocator;
296
297 var queue: MinHeap = .empty;
298 defer queue.deinit(gpa);
299
300 try queue.push(gpa, 54);
301 try queue.push(gpa, 12);
302 try queue.push(gpa, 7);
303 try queue.push(gpa, 23);
304 try queue.push(gpa, 25);
305 try queue.push(gpa, 13);
306 try expectEqual(@as(u32, 7), queue.pop());
307 try expectEqual(@as(u32, 12), queue.pop());
308 try expectEqual(@as(u32, 13), queue.pop());
309 try expectEqual(@as(u32, 23), queue.pop());
310 try expectEqual(@as(u32, 25), queue.pop());
311 try expectEqual(@as(u32, 54), queue.pop());
312}
313
314test "add and remove same min heap" {
315 const gpa = testing.allocator;
316
317 var queue: MinHeap = .empty;
318 defer queue.deinit(gpa);
319
320 try queue.push(gpa, 1);
321 try queue.push(gpa, 1);
322 try queue.push(gpa, 2);
323 try queue.push(gpa, 2);
324 try queue.push(gpa, 1);
325 try queue.push(gpa, 1);
326 try expectEqual(@as(u32, 1), queue.pop());
327 try expectEqual(@as(u32, 1), queue.pop());
328 try expectEqual(@as(u32, 1), queue.pop());
329 try expectEqual(@as(u32, 1), queue.pop());
330 try expectEqual(@as(u32, 2), queue.pop());
331 try expectEqual(@as(u32, 2), queue.pop());
332}
333
334test "removeOrNull on empty" {
335 const gpa = testing.allocator;
336
337 var queue: MinHeap = .empty;
338 defer queue.deinit(gpa);
339
340 try expect(queue.pop() == null);
341}
342
343test "edge case 3 elements" {
344 const gpa = testing.allocator;
345
346 var queue: MinHeap = .empty;
347 defer queue.deinit(gpa);
348
349 try queue.push(gpa, 9);
350 try queue.push(gpa, 3);
351 try queue.push(gpa, 2);
352 try expectEqual(@as(u32, 2), queue.pop());
353 try expectEqual(@as(u32, 3), queue.pop());
354 try expectEqual(@as(u32, 9), queue.pop());
355}
356
357test "peek" {
358 const gpa = testing.allocator;
359
360 var queue: MinHeap = .empty;
361 defer queue.deinit(gpa);
362
363 try expect(queue.peek() == null);
364 try queue.push(gpa, 9);
365 try queue.push(gpa, 3);
366 try queue.push(gpa, 2);
367 try expectEqual(@as(u32, 2), queue.peek().?);
368 try expectEqual(@as(u32, 2), queue.peek().?);
369}
370
371test "sift up with odd indices" {
372 const gpa = testing.allocator;
373
374 var queue: MinHeap = .empty;
375 defer queue.deinit(gpa);
376
377 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
378 for (items) |e| {
379 try queue.push(gpa, e);
380 }
381
382 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
383 for (sorted_items) |e| {
384 try expectEqual(e, queue.pop());
385 }
386}
387
388test "addSlice" {
389 const gpa = testing.allocator;
390
391 var queue: MinHeap = .empty;
392 defer queue.deinit(gpa);
393
394 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
395 try queue.pushSlice(gpa, items[0..]);
396
397 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
398 for (sorted_items) |e| {
399 try expectEqual(e, queue.pop());
400 }
401}
402
403test "fromOwnedSlice trivial case 0" {
404 const gpa = testing.allocator;
405
406 const items = [0]u32{};
407 const queue_items = try gpa.dupe(u32, &items);
408
409 var queue: MinHeap = .fromOwnedSlice(queue_items[0..], {});
410 defer queue.deinit(gpa);
411
412 try expectEqual(@as(usize, 0), queue.count());
413 try expect(queue.pop() == null);
414}
415
416test "fromOwnedSlice trivial case 1" {
417 const gpa = testing.allocator;
418
419 const items = [1]u32{1};
420 const queue_items = try gpa.dupe(u32, &items);
421
422 var queue: MinHeap = .fromOwnedSlice(queue_items[0..], {});
423 defer queue.deinit(gpa);
424
425 try expectEqual(@as(usize, 1), queue.count());
426 try expectEqual(items[0], queue.pop());
427 try expect(queue.pop() == null);
428}
429
430test "fromOwnedSlice" {
431 const gpa = testing.allocator;
432
433 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
434 const heap_items = try gpa.dupe(u32, items[0..]);
435
436 var queue: MinHeap = .fromOwnedSlice(heap_items[0..], {});
437 defer queue.deinit(gpa);
438
439 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
440 for (sorted_items) |e| {
441 try expectEqual(e, queue.pop());
442 }
443}
444
445test "add and remove max heap" {
446 const gpa = testing.allocator;
447
448 var queue: MaxHeap = .empty;
449 defer queue.deinit(gpa);
450
451 try queue.push(gpa, 54);
452 try queue.push(gpa, 12);
453 try queue.push(gpa, 7);
454 try queue.push(gpa, 23);
455 try queue.push(gpa, 25);
456 try queue.push(gpa, 13);
457 try expectEqual(@as(u32, 54), queue.pop());
458 try expectEqual(@as(u32, 25), queue.pop());
459 try expectEqual(@as(u32, 23), queue.pop());
460 try expectEqual(@as(u32, 13), queue.pop());
461 try expectEqual(@as(u32, 12), queue.pop());
462 try expectEqual(@as(u32, 7), queue.pop());
463}
464
465test "add and remove same max heap" {
466 const gpa = testing.allocator;
467
468 var queue: MaxHeap = .empty;
469 defer queue.deinit(gpa);
470
471 try queue.push(gpa, 1);
472 try queue.push(gpa, 1);
473 try queue.push(gpa, 2);
474 try queue.push(gpa, 2);
475 try queue.push(gpa, 1);
476 try queue.push(gpa, 1);
477 try expectEqual(@as(u32, 2), queue.pop());
478 try expectEqual(@as(u32, 2), queue.pop());
479 try expectEqual(@as(u32, 1), queue.pop());
480 try expectEqual(@as(u32, 1), queue.pop());
481 try expectEqual(@as(u32, 1), queue.pop());
482 try expectEqual(@as(u32, 1), queue.pop());
483}
484
485test "iterator" {
486 const gpa = testing.allocator;
487
488 var queue: MinHeap = .empty;
489 var map = std.AutoHashMap(u32, void).init(testing.allocator);
490 defer {
491 queue.deinit(gpa);
492 map.deinit();
493 }
494
495 const items = [_]u32{ 54, 12, 7, 23, 25, 13 };
496 for (items) |e| {
497 _ = try queue.push(gpa, e);
498 try map.put(e, {});
499 }
500
501 var it = queue.iterator();
502 while (it.next()) |e| {
503 _ = map.remove(e);
504 }
505
506 try expectEqual(@as(usize, 0), map.count());
507}
508
509test "remove at index" {
510 const gpa = testing.allocator;
511
512 var queue: MinHeap = .empty;
513 defer queue.deinit(gpa);
514
515 const items = [_]u32{ 2, 1, 8, 9, 3, 4, 5 };
516 for (items) |e| {
517 _ = try queue.push(gpa, e);
518 }
519
520 var it = queue.iterator();
521 var idx: usize = 0;
522 const two_idx = while (it.next()) |elem| {
523 if (elem == 2)
524 break idx;
525 idx += 1;
526 } else unreachable;
527 const sorted_items = [_]u32{ 1, 3, 4, 5, 8, 9 };
528 try expectEqual(queue.popIndex(two_idx), 2);
529
530 var i: usize = 0;
531 while (queue.pop()) |n| : (i += 1) {
532 try expectEqual(n, sorted_items[i]);
533 }
534 try expectEqual(queue.pop(), null);
535}
536
537test "iterator while empty" {
538 const gpa = testing.allocator;
539
540 var queue: MinHeap = .empty;
541 defer queue.deinit(gpa);
542
543 var it = queue.iterator();
544
545 try expectEqual(it.next(), null);
546}
547
548test "shrinkAndFree" {
549 const gpa = testing.allocator;
550
551 var queue: MinHeap = .empty;
552 defer queue.deinit(gpa);
553
554 try queue.ensureTotalCapacity(gpa, 4);
555 try expect(queue.capacity() >= 4);
556
557 try queue.push(gpa, 1);
558 try queue.push(gpa, 2);
559 try queue.push(gpa, 3);
560 try expect(queue.capacity() >= 4);
561 try expectEqual(@as(usize, 3), queue.count());
562
563 queue.shrinkAndFree(gpa, 3);
564 try expectEqual(@as(usize, 3), queue.capacity());
565 try expectEqual(@as(usize, 3), queue.count());
566
567 try expectEqual(@as(u32, 1), queue.pop());
568 try expectEqual(@as(u32, 2), queue.pop());
569 try expectEqual(@as(u32, 3), queue.pop());
570 try expect(queue.pop() == null);
571}
572
573test "update min heap" {
574 const gpa = testing.allocator;
575
576 var queue: MinHeap = .empty;
577 defer queue.deinit(gpa);
578
579 try queue.push(gpa, 55);
580 try queue.push(gpa, 44);
581 try queue.push(gpa, 11);
582 try queue.update(55, 5);
583 try queue.update(44, 4);
584 try queue.update(11, 1);
585 try expectEqual(@as(u32, 1), queue.pop());
586 try expectEqual(@as(u32, 4), queue.pop());
587 try expectEqual(@as(u32, 5), queue.pop());
588}
589
590test "update same min heap" {
591 const gpa = testing.allocator;
592
593 var queue: MinHeap = .empty;
594 defer queue.deinit(gpa);
595
596 try queue.push(gpa, 1);
597 try queue.push(gpa, 1);
598 try queue.push(gpa, 2);
599 try queue.push(gpa, 2);
600 try queue.update(1, 5);
601 try queue.update(2, 4);
602 try expectEqual(@as(u32, 1), queue.pop());
603 try expectEqual(@as(u32, 2), queue.pop());
604 try expectEqual(@as(u32, 4), queue.pop());
605 try expectEqual(@as(u32, 5), queue.pop());
606}
607
608test "update max heap" {
609 const gpa = testing.allocator;
610
611 var queue: MaxHeap = .empty;
612 defer queue.deinit(gpa);
613
614 try queue.push(gpa, 55);
615 try queue.push(gpa, 44);
616 try queue.push(gpa, 11);
617 try queue.update(55, 5);
618 try queue.update(44, 1);
619 try queue.update(11, 4);
620 try expectEqual(@as(u32, 5), queue.pop());
621 try expectEqual(@as(u32, 4), queue.pop());
622 try expectEqual(@as(u32, 1), queue.pop());
623}
624
625test "update same max heap" {
626 const gpa = testing.allocator;
627
628 var queue: MaxHeap = .empty;
629 defer queue.deinit(gpa);
630
631 try queue.push(gpa, 1);
632 try queue.push(gpa, 1);
633 try queue.push(gpa, 2);
634 try queue.push(gpa, 2);
635 try queue.update(1, 5);
636 try queue.update(2, 4);
637 try expectEqual(@as(u32, 5), queue.pop());
638 try expectEqual(@as(u32, 4), queue.pop());
639 try expectEqual(@as(u32, 2), queue.pop());
640 try expectEqual(@as(u32, 1), queue.pop());
641}
642
643test "update after remove" {
644 const gpa = testing.allocator;
645
646 var queue: MinHeap = .empty;
647 defer queue.deinit(gpa);
648
649 try queue.push(gpa, 1);
650 try expectEqual(@as(u32, 1), queue.pop());
651 try expectError(error.ElementNotFound, queue.update(1, 1));
652}
653
654test "siftUp in remove" {
655 const gpa = testing.allocator;
656
657 var queue: MinHeap = .empty;
658 defer queue.deinit(gpa);
659
660 try queue.pushSlice(gpa, &.{ 0, 1, 100, 2, 3, 101, 102, 4, 5, 6, 7, 103, 104, 105, 106, 8 });
661
662 _ = queue.popIndex(std.mem.findScalar(u32, queue.items[0..queue.count()], 102).?);
663
664 const sorted_items = [_]u32{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 100, 101, 103, 104, 105, 106 };
665 for (sorted_items) |e| {
666 try expectEqual(e, queue.pop());
667 }
668}
669
670fn contextLessThan(context: []const u32, a: usize, b: usize) Order {
671 return std.math.order(context[a], context[b]);
672}
673
674const MinHeapWithContext = PriorityQueue(usize, []const u32, contextLessThan);
675
676test "add and remove min heap with context comparator" {
677 const gpa = testing.allocator;
678
679 const context = [_]u32{ 5, 3, 4, 2, 2, 8, 0 };
680
681 var queue: MinHeapWithContext = .initContext(context[0..]);
682 defer queue.deinit(gpa);
683
684 try queue.push(gpa, 0);
685 try queue.push(gpa, 1);
686 try queue.push(gpa, 2);
687 try queue.push(gpa, 3);
688 try queue.push(gpa, 4);
689 try queue.push(gpa, 5);
690 try queue.push(gpa, 6);
691 try expectEqual(@as(usize, 6), queue.pop());
692 try expectEqual(@as(usize, 4), queue.pop());
693 try expectEqual(@as(usize, 3), queue.pop());
694 try expectEqual(@as(usize, 1), queue.pop());
695 try expectEqual(@as(usize, 2), queue.pop());
696 try expectEqual(@as(usize, 0), queue.pop());
697 try expectEqual(@as(usize, 5), queue.pop());
698}