authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2020-10-16 21:26:04-04:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2020-10-16 21:26:04-04:00
log8a2d12d7075fe5eb68e12ccb2b662981b71ae677
tree470a6a6f41677ac3302665b5aa9610b40eb8ae7c
parentd44486b274a916823fcf6045a6400ef51e07d544
parent84406d98e4fd8032b544eba7b5f9d3ed78e67e6c
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

Merge pull request #6393 from onebsv1/priority-queue-update

Adding a function to update the priority of an element

1 files changed, 80 insertions(+), 2 deletions(-)

lib/std/priority_queue.zig+80-2
......@@ -49,7 +49,12 @@ pub fn PriorityQueue(comptime T: type) type {
4949
5050 fn addUnchecked(self: *Self, elem: T) void {
5151 self.items[self.len] = elem;
52 var child_index = self.len;
52 siftUp(self, self.len);
53 self.len += 1;
54 }
55
56 fn siftUp(self: *Self, start_index: usize) void {
57 var child_index = start_index;
5358 while (child_index > 0) {
5459 var parent_index = ((child_index - 1) >> 1);
5560 const child = self.items[child_index];
......@@ -61,7 +66,6 @@ pub fn PriorityQueue(comptime T: type) type {
6166 self.items[child_index] = parent;
6267 child_index = parent_index;
6368 }
64 self.len += 1;
6569 }
6670
6771 /// Add each element in `items` to the queue.
......@@ -190,6 +194,17 @@ pub fn PriorityQueue(comptime T: type) type {
190194 self.len = new_len;
191195 }
192196
197 pub fn update(self: *Self, elem: T, new_elem: T) !void {
198 var update_index: usize = std.mem.indexOfScalar(T, self.items, elem) orelse return error.ElementNotFound;
199 const old_elem: T = self.items[update_index];
200 self.items[update_index] = new_elem;
201 if (self.compareFn(new_elem, old_elem)) {
202 siftUp(self, update_index);
203 } else {
204 siftDown(self, update_index);
205 }
206 }
207
193208 pub const Iterator = struct {
194209 queue: *PriorityQueue(T),
195210 count: usize,
......@@ -437,3 +452,66 @@ test "std.PriorityQueue: iterator while empty" {
437452
438453 expectEqual(it.next(), null);
439454}
455
456test "std.PriorityQueue: update min heap" {
457 var queue = PQ.init(testing.allocator, lessThan);
458 defer queue.deinit();
459
460 try queue.add(55);
461 try queue.add(44);
462 try queue.add(11);
463 try queue.update(55, 5);
464 try queue.update(44, 4);
465 try queue.update(11, 1);
466 expectEqual(@as(u32, 1), queue.remove());
467 expectEqual(@as(u32, 4), queue.remove());
468 expectEqual(@as(u32, 5), queue.remove());
469}
470
471
472test "std.PriorityQueue: update same min heap" {
473 var queue = PQ.init(testing.allocator, lessThan);
474 defer queue.deinit();
475
476 try queue.add(1);
477 try queue.add(1);
478 try queue.add(2);
479 try queue.add(2);
480 try queue.update(1, 5);
481 try queue.update(2, 4);
482 expectEqual(@as(u32, 1), queue.remove());
483 expectEqual(@as(u32, 2), queue.remove());
484 expectEqual(@as(u32, 4), queue.remove());
485 expectEqual(@as(u32, 5), queue.remove());
486}
487
488test "std.PriorityQueue: update max heap" {
489 var queue = PQ.init(testing.allocator, greaterThan);
490 defer queue.deinit();
491
492 try queue.add(55);
493 try queue.add(44);
494 try queue.add(11);
495 try queue.update(55, 5);
496 try queue.update(44, 1);
497 try queue.update(11, 4);
498 expectEqual(@as(u32, 5), queue.remove());
499 expectEqual(@as(u32, 4), queue.remove());
500 expectEqual(@as(u32, 1), queue.remove());
501}
502
503test "std.PriorityQueue: update same max heap" {
504 var queue = PQ.init(testing.allocator, greaterThan);
505 defer queue.deinit();
506
507 try queue.add(1);
508 try queue.add(1);
509 try queue.add(2);
510 try queue.add(2);
511 try queue.update(1, 5);
512 try queue.update(2, 4);
513 expectEqual(@as(u32, 5), queue.remove());
514 expectEqual(@as(u32, 4), queue.remove());
515 expectEqual(@as(u32, 2), queue.remove());
516 expectEqual(@as(u32, 1), queue.remove());
517}
\ No newline at end of file