Commit 7deae07101

joachimschmidt557 <joachim.schmidt557@outlook.com>
2022-05-26 18:23:07
std.PriorityQueue: fix missing siftUp in remove
When the replacement node is smaller than its parent, we need to sift up instead of sifting down.
1 parent b08d32c
Changed files (1)
lib/std/priority_queue.zig
@@ -100,7 +100,19 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
             const item = self.items[index];
             self.items[index] = last;
             self.len -= 1;
-            siftDown(self, index);
+
+            if (index == 0) {
+                siftDown(self, index);
+            } else {
+                const parent_index = ((index - 1) >> 1);
+                const parent = self.items[parent_index];
+                if (compareFn(self.context, last, parent) == .gt) {
+                    siftDown(self, index);
+                } else {
+                    siftUp(self, index);
+                }
+            }
+
             return item;
         }
 
@@ -576,6 +588,20 @@ test "std.PriorityQueue: update same max heap" {
     try expectEqual(@as(u32, 1), queue.remove());
 }
 
+test "std.PriorityQueue: siftUp in remove" {
+    var queue = PQlt.init(testing.allocator, {});
+    defer queue.deinit();
+
+    try queue.addSlice(&.{ 0, 1, 100, 2, 3, 101, 102, 4, 5, 6, 7, 103, 104, 105, 106, 8 });
+
+    _ = queue.removeIndex(std.mem.indexOfScalar(u32, queue.items[0..queue.len], 102).?);
+
+    const sorted_items = [_]u32{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 100, 101, 103, 104, 105, 106 };
+    for (sorted_items) |e| {
+        try expectEqual(e, queue.remove());
+    }
+}
+
 fn contextLessThan(context: []const u32, a: usize, b: usize) Order {
     return std.math.order(context[a], context[b]);
 }