Priority Dequeue for storing generic data. Initialize with init.
Provide compareFn that returns Order.lt when its second
argument should get min-popped before its third argument,
Order.eq if the arguments are of equal priority, or Order.gt
if the third argument should be min-popped second.
Popping the max element works in reverse. For example,
to make popMin return the smallest number, provide
fn lessThan(context: void, a: T, b: T) Order { _ = context; return std.math.order(a, b); }
pub fn PriorityDequeue(comptime T: type, comptime Context: type, comptime compareFn: fn (context: Context, a: T, b: T) Order) typepub fn PriorityDequeue(comptime T: type, comptime Context: type, comptime compareFn: fn (context: Context, a: T, b: T) Order) type {
return struct {
const Self = @This();
items: []T,
len: usize,
context: Context,
/// A priority dequeue containing no elements.
pub const empty: Self = .{
.items = &.{},
.len = 0,
.context = undefined,
};
/// Initialize and return a new priority dequeue with context.
pub fn initContext(context: Context) Self {
return Self{
.items = &.{},
.len = 0,
.context = context,
};
}
/// Free memory used by the dequeue.
pub fn deinit(self: *Self, allocator: Allocator) void {
allocator.free(self.items);
self.* = undefined;
}
/// Insert a new element, maintaining priority.
pub fn push(self: *Self, allocator: Allocator, elem: T) !void {
try self.ensureUnusedCapacity(allocator, 1);
pushUnchecked(self, elem);
}
/// Add each element in `items` to the dequeue.
pub fn pushSlice(self: *Self, allocator: Allocator, items: []const T) !void {
try self.ensureUnusedCapacity(allocator, items.len);
for (items) |e| {
self.pushUnchecked(e);
}
}
fn pushUnchecked(self: *Self, elem: T) void {
self.items[self.len] = elem;
if (self.len > 0) {
const start = self.getStartForSiftUp(elem, self.len);
self.siftUp(start);
}
self.len += 1;
}
fn isMinLayer(index: usize) bool {
// In the min-max heap structure:
// The first element is on a min layer;
// next two are on a max layer;
// next four are on a min layer, and so on.
return 1 == @clz(index +% 1) & 1;
}
fn nextIsMinLayer(self: *const Self) bool {
return isMinLayer(self.len);
}
const StartIndexAndLayer = struct {
index: usize,
min_layer: bool,
};
fn getStartForSiftUp(self: *const Self, child: T, index: usize) StartIndexAndLayer {
const child_index = index;
const parent_index = parentIndex(child_index);
const parent = self.items[parent_index];
const min_layer = self.nextIsMinLayer();
const order = compareFn(self.context, child, parent);
if ((min_layer and order == .gt) or (!min_layer and order == .lt)) {
// We must swap the item with it's parent if it is on the "wrong" layer
self.items[parent_index] = child;
self.items[child_index] = parent;
return .{
.index = parent_index,
.min_layer = !min_layer,
};
} else {
return .{
.index = child_index,
.min_layer = min_layer,
};
}
}
fn siftUp(self: *Self, start: StartIndexAndLayer) void {
if (start.min_layer) {
doSiftUp(self, start.index, .lt);
} else {
doSiftUp(self, start.index, .gt);
}
}
fn doSiftUp(self: *Self, start_index: usize, target_order: Order) void {
var child_index = start_index;
while (child_index > 2) {
const grandparent_index = grandparentIndex(child_index);
const child = self.items[child_index];
const grandparent = self.items[grandparent_index];
// If the grandparent is already better or equal, we have gone as far as we need to
if (compareFn(self.context, child, grandparent) != target_order) break;
// Otherwise swap the item with it's grandparent
self.items[grandparent_index] = child;
self.items[child_index] = grandparent;
child_index = grandparent_index;
}
}
/// Look at the smallest element in the dequeue. Returns
/// `null` if empty.
pub fn peekMin(self: *const Self) ?T {
return if (self.len > 0) self.items[0] else null;
}
/// Look at the largest element in the dequeue. Returns
/// `null` if empty.
pub fn peekMax(self: *const Self) ?T {
if (self.len == 0) return null;
if (self.len == 1) return self.items[0];
if (self.len == 2) return self.items[1];
return self.bestItemAtIndices(1, 2, .gt).item;
}
fn maxIndex(self: *const Self) ?usize {
if (self.len == 0) return null;
if (self.len == 1) return 0;
if (self.len == 2) return 1;
return self.bestItemAtIndices(1, 2, .gt).index;
}
/// Remove and return the smallest element from the dequeue, or `null` if empty
pub fn popMin(self: *Self) ?T {
return if (self.len > 0) self.popIndex(0) else null;
}
/// Remove and return the largest element from the dequeue, or `null` if empty
pub fn popMax(self: *Self) ?T {
return if (self.len > 0) self.popIndex(self.maxIndex().?) else null;
}
/// Remove and return element at index. Indices are in the
/// same order as iterator, which is not necessarily priority
/// order.
pub fn popIndex(self: *Self, index: usize) T {
assert(self.len > index);
const item = self.items[index];
const last = self.items[self.len - 1];
self.items[index] = last;
self.len -= 1;
siftDown(self, index);
return item;
}
fn siftDown(self: *Self, index: usize) void {
if (isMinLayer(index)) {
self.doSiftDown(index, .lt);
} else {
self.doSiftDown(index, .gt);
}
}
fn doSiftDown(self: *Self, start_index: usize, target_order: Order) void {
var index = start_index;
const half = self.len >> 1;
while (true) {
const first_grandchild_index = firstGrandchildIndex(index);
const last_grandchild_index = first_grandchild_index + 3;
const elem = self.items[index];
if (last_grandchild_index < self.len) {
// All four grandchildren exist
const index2 = first_grandchild_index + 1;
const index3 = index2 + 1;
// Find the best grandchild
const best_left = self.bestItemAtIndices(first_grandchild_index, index2, target_order);
const best_right = self.bestItemAtIndices(index3, last_grandchild_index, target_order);
const best_grandchild = self.bestItem(best_left, best_right, target_order);
// If the item is better than or equal to its best grandchild, we are done
if (compareFn(self.context, best_grandchild.item, elem) != target_order) return;
// Otherwise, swap them
self.items[best_grandchild.index] = elem;
self.items[index] = best_grandchild.item;
index = best_grandchild.index;
// We might need to swap the element with it's parent
self.swapIfParentIsBetter(elem, index, target_order);
} else {
// The children or grandchildren are the last layer
const first_child_index = firstChildIndex(index);
if (first_child_index >= self.len) return;
const best_descendent = self.bestDescendent(first_child_index, first_grandchild_index, target_order);
// If the item is better than or equal to its best descendant, we are done
if (compareFn(self.context, best_descendent.item, elem) != target_order) return;
// Otherwise swap them
self.items[best_descendent.index] = elem;
self.items[index] = best_descendent.item;
index = best_descendent.index;
// If we didn't swap a grandchild, we are done
if (index < first_grandchild_index) return;
// We might need to swap the element with it's parent
self.swapIfParentIsBetter(elem, index, target_order);
return;
}
// If we are now in the last layer, we are done
if (index >= half) return;
}
}
fn swapIfParentIsBetter(self: *Self, child: T, child_index: usize, target_order: Order) void {
const parent_index = parentIndex(child_index);
const parent = self.items[parent_index];
if (compareFn(self.context, parent, child) == target_order) {
self.items[parent_index] = child;
self.items[child_index] = parent;
}
}
const ItemAndIndex = struct {
item: T,
index: usize,
};
fn getItem(self: *const Self, index: usize) ItemAndIndex {
return .{
.item = self.items[index],
.index = index,
};
}
fn bestItem(self: *const Self, item1: ItemAndIndex, item2: ItemAndIndex, target_order: Order) ItemAndIndex {
if (compareFn(self.context, item1.item, item2.item) == target_order) {
return item1;
} else {
return item2;
}
}
fn bestItemAtIndices(self: *const Self, index1: usize, index2: usize, target_order: Order) ItemAndIndex {
const item1 = self.getItem(index1);
const item2 = self.getItem(index2);
return self.bestItem(item1, item2, target_order);
}
fn bestDescendent(self: *const Self, first_child_index: usize, first_grandchild_index: usize, target_order: Order) ItemAndIndex {
const second_child_index = first_child_index + 1;
if (first_grandchild_index >= self.len) {
// No grandchildren, find the best child (second may not exist)
if (second_child_index >= self.len) {
return .{
.item = self.items[first_child_index],
.index = first_child_index,
};
} else {
return self.bestItemAtIndices(first_child_index, second_child_index, target_order);
}
}
const second_grandchild_index = first_grandchild_index + 1;
if (second_grandchild_index >= self.len) {
// One grandchild, so we know there is a second child. Compare first grandchild and second child
return self.bestItemAtIndices(first_grandchild_index, second_child_index, target_order);
}
const best_left_grandchild_index = self.bestItemAtIndices(first_grandchild_index, second_grandchild_index, target_order).index;
const third_grandchild_index = second_grandchild_index + 1;
if (third_grandchild_index >= self.len) {
// Two grandchildren, and we know the best. Compare this to second child.
return self.bestItemAtIndices(best_left_grandchild_index, second_child_index, target_order);
} else {
// Three grandchildren, compare the min of the first two with the third
return self.bestItemAtIndices(best_left_grandchild_index, third_grandchild_index, target_order);
}
}
/// Return the number of elements remaining in the dequeue
pub fn count(self: *const Self) usize {
return self.len;
}
/// Return the number of elements that can be added to the
/// dequeue before more memory is allocated.
pub fn capacity(self: *const Self) usize {
return self.items.len;
}
/// Dequeue takes ownership of the passed in slice. The slice must be de-initialize
/// with `deinit`.
pub fn fromOwnedSlice(items: []T, context: Context) Self {
var queue = Self{
.items = items,
.len = items.len,
.context = context,
};
if (queue.len <= 1) return queue;
const half = (queue.len >> 1) - 1;
var i: usize = 0;
while (i <= half) : (i += 1) {
const index = half - i;
queue.siftDown(index);
}
return queue;
}
/// Ensure that the dequeue can fit at least `new_capacity` items.
pub fn ensureTotalCapacity(self: *Self, allocator: Allocator, new_capacity: usize) !void {
var better_capacity = self.capacity();
if (better_capacity >= new_capacity) return;
while (true) {
better_capacity += better_capacity / 2 + 8;
if (better_capacity >= new_capacity) break;
}
self.items = try allocator.realloc(self.items, better_capacity);
}
/// Ensure that the dequeue can fit at least `additional_count` **more** items.
pub fn ensureUnusedCapacity(self: *Self, allocator: Allocator, additional_count: usize) !void {
return self.ensureTotalCapacity(allocator, self.len + additional_count);
}
/// Reduce allocated capacity to `new_len`.
pub fn shrinkAndFree(self: *Self, allocator: Allocator, new_len: usize) void {
assert(new_len <= self.items.len);
// Cannot shrink to smaller than the current queue size without invalidating the heap property
assert(new_len >= self.len);
self.items = allocator.realloc(self.items[0..], new_len) catch |e| switch (e) {
error.OutOfMemory => { // no problem, capacity is still correct then.
self.items.len = new_len;
return;
},
};
}
pub fn update(self: *Self, elem: T, new_elem: T) !void {
const old_index = blk: {
var idx: usize = 0;
while (idx < self.len) : (idx += 1) {
const item = self.items[idx];
if (compareFn(self.context, item, elem) == .eq) break :blk idx;
}
return error.ElementNotFound;
};
_ = self.popIndex(old_index);
self.pushUnchecked(new_elem);
}
pub const Iterator = struct {
queue: *PriorityDequeue(T, Context, compareFn),
count: usize,
pub fn next(it: *Iterator) ?T {
if (it.count >= it.queue.len) return null;
const out = it.count;
it.count += 1;
return it.queue.items[out];
}
pub fn reset(it: *Iterator) void {
it.count = 0;
}
};
/// Return an iterator that walks the queue without consuming
/// it. The iteration order may differ from the priority order.
/// Invalidated if the queue is modified.
pub fn iterator(self: *Self) Iterator {
return Iterator{
.queue = self,
.count = 0,
};
}
fn dump(self: *Self) void {
const print = std.debug.print;
print("{{ ", .{});
print("items: ", .{});
for (self.items, 0..) |e, i| {
if (i >= self.len) break;
print("{}, ", .{e});
}
print("array: ", .{});
for (self.items) |e| {
print("{}, ", .{e});
}
print("len: {} ", .{self.len});
print("capacity: {}", .{self.capacity()});
print(" }}\n", .{});
}
fn parentIndex(index: usize) usize {
return (index - 1) >> 1;
}
fn grandparentIndex(index: usize) usize {
return parentIndex(parentIndex(index));
}
fn firstChildIndex(index: usize) usize {
return (index << 1) + 1;
}
fn firstGrandchildIndex(index: usize) usize {
return firstChildIndex(firstChildIndex(index));
}
};
}