Commit 7deae07101
Changed files (1)
lib
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]);
}