feat: add MinHeap and deinit
This commit is contained in:
1 parent
73c82254f9
commit
1b7b3fd744
1 file changed
+156
+156
@@ -0,0 +1,156 @@
|
|||||||
|
const std = @import("std");
|
||||||
|
const Allocator = std.mem.Allocator;
|
||||||
|
|
||||||
|
/// Binary heap ordered by `lessThan(context, a, b)`: parent never compares
|
||||||
|
/// greater than either child under that relation. Pass a comparator that
|
||||||
|
/// returns `a > b` to get a max-heap out of the same implementation.
|
||||||
|
pub fn MinHeap(comptime T: type, comptime Context: type, comptime lessThan: fn (context: Context, a: T, b: T) bool) type {
|
||||||
|
return struct {
|
||||||
|
items: std.ArrayList(T),
|
||||||
|
context: Context,
|
||||||
|
|
||||||
|
const Self = @This();
|
||||||
|
|
||||||
|
pub fn init(allocator: Allocator, context: Context) Self {
|
||||||
|
return .{ .items = std.ArrayList(T).init(allocator), .context = context };
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn deinit(self: *Self) void {
|
||||||
|
self.items.deinit();
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn len(self: *const Self) usize {
|
||||||
|
return self.items.items.len;
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn peek(self: *const Self) ?T {
|
||||||
|
if (self.items.items.len == 0) return null;
|
||||||
|
return self.items.items[0];
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn push(self: *Self, item: T) !void {
|
||||||
|
try self.items.append(item);
|
||||||
|
self.siftUp(self.items.items.len - 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn pop(self: *Self) ?T {
|
||||||
|
if (self.items.items.len == 0) return null;
|
||||||
|
const top = self.items.items[0];
|
||||||
|
const last = self.items.pop();
|
||||||
|
if (self.items.items.len > 0) {
|
||||||
|
self.items.items[0] = last;
|
||||||
|
self.siftDown(0);
|
||||||
|
}
|
||||||
|
return top;
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Swap the root for `item` and restore the heap property in one
|
||||||
|
/// siftDown, instead of a pop() (siftDown) followed by a push()
|
||||||
|
/// (siftUp). Caller must ensure the heap is non-empty.
|
||||||
|
pub fn replaceTop(self: *Self, item: T) T {
|
||||||
|
const old = self.items.items[0];
|
||||||
|
self.items.items[0] = item;
|
||||||
|
self.siftDown(0);
|
||||||
|
return old;
|
||||||
|
}
|
||||||
|
|
||||||
|
fn siftUp(self: *Self, startIdx: usize) void {
|
||||||
|
var idx = startIdx;
|
||||||
|
while (idx > 0) {
|
||||||
|
const parent = (idx - 1) / 2;
|
||||||
|
if (!lessThan(self.context, self.items.items[idx], self.items.items[parent])) break;
|
||||||
|
std.mem.swap(T, &self.items.items[idx], &self.items.items[parent]);
|
||||||
|
idx = parent;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn siftDown(self: *Self, startIdx: usize) void {
|
||||||
|
var idx = startIdx;
|
||||||
|
const n = self.items.items.len;
|
||||||
|
while (true) {
|
||||||
|
const left = idx * 2 + 1;
|
||||||
|
const right = idx * 2 + 2;
|
||||||
|
var smallest = idx;
|
||||||
|
if (left < n and lessThan(self.context, self.items.items[left], self.items.items[smallest])) smallest = left;
|
||||||
|
if (right < n and lessThan(self.context, self.items.items[right], self.items.items[smallest])) smallest = right;
|
||||||
|
if (smallest == idx) break;
|
||||||
|
std.mem.swap(T, &self.items.items[idx], &self.items.items[smallest]);
|
||||||
|
idx = smallest;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Debug helper: verifies the heap invariant holds at every node.
|
||||||
|
pub fn isValid(self: *const Self) bool {
|
||||||
|
const n = self.items.items.len;
|
||||||
|
var idx: usize = 0;
|
||||||
|
while (idx < n) : (idx += 1) {
|
||||||
|
const left = idx * 2 + 1;
|
||||||
|
const right = idx * 2 + 2;
|
||||||
|
if (left < n and lessThan(self.context, self.items.items[left], self.items.items[idx])) return false;
|
||||||
|
if (right < n and lessThan(self.context, self.items.items[right], self.items.items[idx])) return false;
|
||||||
|
}
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
fn u32Less(_: void, a: u32, b: u32) bool {
|
||||||
|
return a < b;
|
||||||
|
}
|
||||||
|
const U32Heap = MinHeap(u32, void, u32Less);
|
||||||
|
|
||||||
|
test "heap: empty peek/pop return null" {
|
||||||
|
var heap = U32Heap.init(std.testing.allocator, {});
|
||||||
|
defer heap.deinit();
|
||||||
|
try std.testing.expectEqual(@as(?u32, null), heap.peek());
|
||||||
|
try std.testing.expectEqual(@as(?u32, null), heap.pop());
|
||||||
|
try std.testing.expectEqual(@as(usize, 0), heap.len());
|
||||||
|
}
|
||||||
|
|
||||||
|
test "heap: push/pop yields ascending order" {
|
||||||
|
var heap = U32Heap.init(std.testing.allocator, {});
|
||||||
|
defer heap.deinit();
|
||||||
|
const values = [_]u32{ 5, 1, 9, 3, 3, 7, 0, 42, 2 };
|
||||||
|
for (values) |v| try heap.push(v);
|
||||||
|
try std.testing.expectEqual(values.len, heap.len());
|
||||||
|
|
||||||
|
var sorted = values;
|
||||||
|
std.mem.sort(u32, &sorted, {}, std.sort.asc(u32));
|
||||||
|
|
||||||
|
for (sorted) |expected| {
|
||||||
|
try std.testing.expectEqual(@as(?u32, expected), heap.peek());
|
||||||
|
try std.testing.expectEqual(@as(?u32, expected), heap.pop());
|
||||||
|
}
|
||||||
|
try std.testing.expectEqual(@as(usize, 0), heap.len());
|
||||||
|
}
|
||||||
|
|
||||||
|
test "heap: peek does not remove" {
|
||||||
|
var heap = U32Heap.init(std.testing.allocator, {});
|
||||||
|
defer heap.deinit();
|
||||||
|
try heap.push(10);
|
||||||
|
try heap.push(4);
|
||||||
|
try std.testing.expectEqual(@as(?u32, 4), heap.peek());
|
||||||
|
try std.testing.expectEqual(@as(?u32, 4), heap.peek());
|
||||||
|
try std.testing.expectEqual(@as(usize, 2), heap.len());
|
||||||
|
}
|
||||||
|
|
||||||
|
test "heap: replaceTop matches pop-then-push" {
|
||||||
|
var a = U32Heap.init(std.testing.allocator, {});
|
||||||
|
defer a.deinit();
|
||||||
|
var b = U32Heap.init(std.testing.allocator, {});
|
||||||
|
defer b.deinit();
|
||||||
|
for ([_]u32{ 8, 2, 6, 1, 9 }) |v| {
|
||||||
|
try a.push(v);
|
||||||
|
try b.push(v);
|
||||||
|
}
|
||||||
|
|
||||||
|
const old = a.replaceTop(5);
|
||||||
|
try std.testing.expectEqual(@as(u32, 1), old);
|
||||||
|
_ = b.pop();
|
||||||
|
try b.push(5);
|
||||||
|
|
||||||
|
while (a.len() > 0) {
|
||||||
|
try std.testing.expectEqual(a.pop().?, b.pop().?);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
Reference in new issue
Block a user