| ... | ... | @@ -1,8 +1,10 @@ |
| 1 | 1 | const std = @import("std.zig"); |
| 2 | 2 | const Allocator = std.mem.Allocator; |
| 3 | 3 | const debug = std.debug; |
| 4 | const assert = debug.assert; |
| 4 | 5 | const expect = std.testing.expect; |
| 5 | 6 | const expectEqual = std.testing.expectEqual; |
| 7 | const expectError = std.testing.expectError; |
| 6 | 8 | |
| 7 | 9 | /// Priority queue for storing generic data. Initialize with `init`. |
| 8 | 10 | pub fn PriorityQueue(comptime T: type) type { |
| ... | ... | @@ -77,13 +79,23 @@ pub fn PriorityQueue(comptime T: type) type { |
| 77 | 79 | return if (self.len > 0) self.remove() else null; |
| 78 | 80 | } |
| 79 | 81 | |
| 82 | /// Remove and return the highest priority element from the |
| 83 | /// queue. |
| 80 | 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 | 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 | 96 | self.len -= 1; |
| 85 | 97 | siftDown(self, 0); |
| 86 | | return first; |
| 98 | return item; |
| 87 | 99 | } |
| 88 | 100 | |
| 89 | 101 | /// Return the number of elements remaining in the priority |
| ... | ... | @@ -388,3 +400,26 @@ test "std.PriorityQueue: iterator" { |
| 388 | 400 | |
| 389 | 401 | expectEqual(@as(usize, 0), map.count()); |
| 390 | 402 | } |
| 403 | |
| 404 | test "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 | } |
Nathan Michaels <nathan@nmichaels.org> 2020-01-08 13:55:47-05:00
Andrew Kelley <andrew@ziglang.org> 2020-01-08 13:55:47-05:00
lib