authorgravatar for nathan@nmichaels.orgNathan Michaels <nathan@nmichaels.org> 2020-01-08 13:55:47-05:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2020-01-08 13:55:47-05:00
log38ce7f64e3fd81476d7fda1dd85d10e493a8ee04
tree1789d2d37a270a6fa7f206e5a789245c747c9d09
parent0ebb07f95d4658f6ac54b381a049ab0ad79442ab

Add removeIndex function to PriorityQueue (#4070)

It's awkward to use, but lets me cancel events in an event queue. Co-authored-by: Dmitry Atamanov <data-man@users.noreply.github.com>

1 files changed, 38 insertions(+), 3 deletions(-)

lib/std/priority_queue.zig+38-3
...@@ -1,8 +1,10 @@...@@ -1,8 +1,10 @@
1const std = @import("std.zig");1const std = @import("std.zig");
2const Allocator = std.mem.Allocator;2const Allocator = std.mem.Allocator;
3const debug = std.debug;3const debug = std.debug;
4const assert = debug.assert;
4const expect = std.testing.expect;5const expect = std.testing.expect;
5const expectEqual = std.testing.expectEqual;6const expectEqual = std.testing.expectEqual;
7const expectError = std.testing.expectError;
68
7/// Priority queue for storing generic data. Initialize with `init`.9/// Priority queue for storing generic data. Initialize with `init`.
8pub fn PriorityQueue(comptime T: type) type {10pub fn PriorityQueue(comptime T: type) type {
...@@ -77,13 +79,23 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -77,13 +79,23 @@ pub fn PriorityQueue(comptime T: type) type {
77 return if (self.len > 0) self.remove() else null;79 return if (self.len > 0) self.remove() else null;
78 }80 }
7981
82 /// Remove and return the highest priority element from the
83 /// queue.
80 pub fn remove(self: *Self) T {84 pub fn remove(self: *Self) T {
81 const first = self.items[0];85 return self.removeIndex(0);
86 }
87
88 /// Remove and return element at index. Indices are in the
89 /// same order as iterator, which is not necessarily priority
90 /// order.
91 pub fn removeIndex(self: *Self, index: usize) T {
92 assert(self.len > index);
82 const last = self.items[self.len - 1];93 const last = self.items[self.len - 1];
83 self.items[0] = last;94 const item = self.items[index];
95 self.items[index] = last;
84 self.len -= 1;96 self.len -= 1;
85 siftDown(self, 0);97 siftDown(self, 0);
86 return first;98 return item;
87 }99 }
88100
89 /// Return the number of elements remaining in the priority101 /// Return the number of elements remaining in the priority
...@@ -388,3 +400,26 @@ test "std.PriorityQueue: iterator" {...@@ -388,3 +400,26 @@ test "std.PriorityQueue: iterator" {
388400
389 expectEqual(@as(usize, 0), map.count());401 expectEqual(@as(usize, 0), map.count());
390}402}
403
404test "std.PriorityQueue: remove at index" {
405 var queue = PQ.init(debug.global_allocator, lessThan);
406 defer queue.deinit();
407
408 try queue.add(3);
409 try queue.add(2);
410 try queue.add(1);
411
412 var it = queue.iterator();
413 var elem = it.next();
414 var idx: usize = 0;
415 const two_idx = while (elem != null) : (elem = it.next()) {
416 if (elem.? == 2)
417 break idx;
418 idx += 1;
419 } else unreachable;
420
421 expectEqual(queue.removeIndex(two_idx), 2);
422 expectEqual(queue.remove(), 1);
423 expectEqual(queue.remove(), 3);
424 expectEqual(queue.removeOrNull(), null);
425}