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.
This commit is contained in:
joachimschmidt557
2022-05-26 18:23:07 +02:00
committed by Andrew Kelley
parent b08d32ceb5
commit 7deae07101

View File

@@ -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]);
}