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>
master
Nathan Michaels 2020-01-08 13:55:47 -05:00 committed by Andrew Kelley
parent 0ebb07f95d
commit 38ce7f64e3
1 changed files with 38 additions and 3 deletions

View File

@ -1,8 +1,10 @@
const std = @import("std.zig");
const Allocator = std.mem.Allocator;
const debug = std.debug;
const assert = debug.assert;
const expect = std.testing.expect;
const expectEqual = std.testing.expectEqual;
const expectError = std.testing.expectError;
/// Priority queue for storing generic data. Initialize with `init`.
pub fn PriorityQueue(comptime T: type) type {
@ -77,13 +79,23 @@ pub fn PriorityQueue(comptime T: type) type {
return if (self.len > 0) self.remove() else null;
}
/// Remove and return the highest priority element from the
/// queue.
pub fn remove(self: *Self) T {
const first = self.items[0];
return self.removeIndex(0);
}
/// Remove and return element at index. Indices are in the
/// same order as iterator, which is not necessarily priority
/// order.
pub fn removeIndex(self: *Self, index: usize) T {
assert(self.len > index);
const last = self.items[self.len - 1];
self.items[0] = last;
const item = self.items[index];
self.items[index] = last;
self.len -= 1;
siftDown(self, 0);
return first;
return item;
}
/// Return the number of elements remaining in the priority
@ -388,3 +400,26 @@ test "std.PriorityQueue: iterator" {
expectEqual(@as(usize, 0), map.count());
}
test "std.PriorityQueue: remove at index" {
var queue = PQ.init(debug.global_allocator, lessThan);
defer queue.deinit();
try queue.add(3);
try queue.add(2);
try queue.add(1);
var it = queue.iterator();
var elem = it.next();
var idx: usize = 0;
const two_idx = while (elem != null) : (elem = it.next()) {
if (elem.? == 2)
break idx;
idx += 1;
} else unreachable;
expectEqual(queue.removeIndex(two_idx), 2);
expectEqual(queue.remove(), 1);
expectEqual(queue.remove(), 3);
expectEqual(queue.removeOrNull(), null);
}