diff --git a/lib/std/priority_queue.zig b/lib/std/priority_queue.zig index 52d9f59951..ebc13a9974 100644 --- a/lib/std/priority_queue.zig +++ b/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]); }