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

diff --git a/lib/std/priority_queue.zig b/lib/std/priority_queue.zig
index b9be9b70bf..f5ff4800e8 100644
--- a/lib/std/priority_queue.zig
+++ b/lib/std/priority_queue.zig
@@ -49,7 +49,12 @@ pub fn PriorityQueue(comptime T: type) type {
 
         fn addUnchecked(self: *Self, elem: T) void {
             self.items[self.len] = elem;
-            var child_index = self.len;
+            siftUp(self, self.len);
+            self.len += 1;
+        }
+
+        fn siftUp(self: *Self, start_index: usize) void {
+            var child_index = start_index;
             while (child_index > 0) {
                 var parent_index = ((child_index - 1) >> 1);
                 const child = self.items[child_index];
@@ -61,7 +66,6 @@ pub fn PriorityQueue(comptime T: type) type {
                 self.items[child_index] = parent;
                 child_index = parent_index;
             }
-            self.len += 1;
         }
 
         /// Add each element in `items` to the queue.
@@ -190,6 +194,17 @@ pub fn PriorityQueue(comptime T: type) type {
             self.len = new_len;
         }
 
+        pub fn update(self: *Self, elem: T, new_elem: T) !void {
+            var update_index: usize = std.mem.indexOfScalar(T, self.items, elem) orelse return error.ElementNotFound;
+            const old_elem: T = self.items[update_index];
+            self.items[update_index] = new_elem;
+            if (self.compareFn(new_elem, old_elem)) {
+                siftUp(self, update_index);
+            } else {
+                siftDown(self, update_index);
+            }
+        }
+
         pub const Iterator = struct {
             queue: *PriorityQueue(T),
             count: usize,
@@ -437,3 +452,66 @@ test "std.PriorityQueue: iterator while empty" {
 
     expectEqual(it.next(), null);
 }
+
+test "std.PriorityQueue: update min heap" {
+    var queue = PQ.init(testing.allocator, lessThan);
+    defer queue.deinit();
+
+    try queue.add(55);
+    try queue.add(44);
+    try queue.add(11);
+    try queue.update(55, 5);
+    try queue.update(44, 4);
+    try queue.update(11, 1);
+    expectEqual(@as(u32, 1), queue.remove());
+    expectEqual(@as(u32, 4), queue.remove());
+    expectEqual(@as(u32, 5), queue.remove());
+}
+
+
+test "std.PriorityQueue: update same min heap" {
+    var queue = PQ.init(testing.allocator, lessThan);
+    defer queue.deinit();
+
+    try queue.add(1);
+    try queue.add(1);
+    try queue.add(2);
+    try queue.add(2);
+    try queue.update(1, 5);
+    try queue.update(2, 4);
+    expectEqual(@as(u32, 1), queue.remove());
+    expectEqual(@as(u32, 2), queue.remove());
+    expectEqual(@as(u32, 4), queue.remove());
+    expectEqual(@as(u32, 5), queue.remove());
+}
+
+test "std.PriorityQueue: update max heap" {
+    var queue = PQ.init(testing.allocator, greaterThan);
+    defer queue.deinit();
+
+    try queue.add(55);
+    try queue.add(44);
+    try queue.add(11);
+    try queue.update(55, 5);
+    try queue.update(44, 1);
+    try queue.update(11, 4);
+    expectEqual(@as(u32, 5), queue.remove());
+    expectEqual(@as(u32, 4), queue.remove());
+    expectEqual(@as(u32, 1), queue.remove());
+}
+
+test "std.PriorityQueue: update same max heap" {
+    var queue = PQ.init(testing.allocator, greaterThan);
+    defer queue.deinit();
+
+    try queue.add(1);
+    try queue.add(1);
+    try queue.add(2);
+    try queue.add(2);
+    try queue.update(1, 5);
+    try queue.update(2, 4);
+    expectEqual(@as(u32, 5), queue.remove());
+    expectEqual(@as(u32, 4), queue.remove());
+    expectEqual(@as(u32, 2), queue.remove());
+    expectEqual(@as(u32, 1), queue.remove());
+}
\ No newline at end of file