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 {...@@ -49,7 +49,12 @@ pub fn PriorityQueue(comptime T: type) type {
4949
50 fn addUnchecked(self: *Self, elem: T) void {50 fn addUnchecked(self: *Self, elem: T) void {
51 self.items[self.len] = elem;51 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;
53 while (child_index > 0) {58 while (child_index > 0) {
54 var parent_index = ((child_index - 1) >> 1);59 var parent_index = ((child_index - 1) >> 1);
55 const child = self.items[child_index];60 const child = self.items[child_index];
...@@ -61,7 +66,6 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -61,7 +66,6 @@ pub fn PriorityQueue(comptime T: type) type {
61 self.items[child_index] = parent;66 self.items[child_index] = parent;
62 child_index = parent_index;67 child_index = parent_index;
63 }68 }
64 self.len += 1;
65 }69 }
6670
67 /// Add each element in `items` to the queue.71 /// Add each element in `items` to the queue.
...@@ -190,6 +194,17 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -190,6 +194,17 @@ pub fn PriorityQueue(comptime T: type) type {
190 self.len = new_len;194 self.len = new_len;
191 }195 }
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
193 pub const Iterator = struct {208 pub const Iterator = struct {
194 queue: *PriorityQueue(T),209 queue: *PriorityQueue(T),
195 count: usize,210 count: usize,
...@@ -437,3 +452,66 @@ test "std.PriorityQueue: iterator while empty" {...@@ -437,3 +452,66 @@ test "std.PriorityQueue: iterator while empty" {
437452
438 expectEqual(it.next(), null);453 expectEqual(it.next(), null);
439}454}
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