priorityq/binheap/lib.go

// Package binheap implements a binary max-heap.
package binheap

import "golang.org/x/exp/constraints"

// H is a binary max-heap.
//
// `I` is the type of the priority IDs, and `E` the type of the elements.
type H[I constraints.Ordered, E any] struct {
	heap  []I
	elems []E
	len   int
}

// Make creates a new heap.
func Make[I constraints.Ordered, E any](cap int) H[I, E] {
	heap := make([]I, cap)
	elems := make([]E, cap)
	h := H[I, E]{heap: heap, elems: elems}
	return h
}

// Capacity returns the total capacity of the heap.
func (h *H[I, E]) Capacity() int {
	return cap(h.heap)
}

// Len returns the number of items in the heap.
func (h *H[I, E]) Len() int {
	return h.len
}

// CanExtract returns true if the heap has any item, otherwise false.
func (h *H[I, E]) CanExtract() bool {
	return h.len != 0
}

// CanInsert returns true if the heap has unused capacity, otherwise false.
func (h *H[I, E]) CanInsert() bool {
	return cap(h.heap)-h.len != 0
}

// Extract returns the current heap root, then performs a heap-down pass.
//
// If the heap is empty, it panics.
func (h *H[I, E]) Extract() (I, E) {
	if !h.CanExtract() {
		panic("heap is empty")
	}

	id := h.heap[0]
	elem := h.elems[0]
	var emptyId I
	var emptyElem E
	h.heap[0] = h.heap[h.len-1]
	h.elems[0] = h.elems[h.len-1]
	h.heap[h.len-1] = emptyId
	h.elems[h.len-1] = emptyElem
	h.len--
	idx := 0
	for {
		left := idx*2 + 1
		right := idx*2 + 2
		largest := idx
		if left < h.len && h.heap[left] > h.heap[largest] {
			largest = left
		}
		if right < h.len && h.heap[right] > h.heap[largest] {
			largest = right
		}
		if largest == idx {
			break
		}
		h.heap[idx], h.heap[largest] = h.heap[largest], h.heap[idx]
		h.elems[idx], h.elems[largest] = h.elems[largest], h.elems[idx]
		idx = largest
	}

	return id, elem
}

// Insert adds an item to the heap, then performs a heap-up pass.
//
// If the heap is full, it panics.
func (h *H[I, E]) Insert(id I, elem E) {
	if !h.CanInsert() {
		panic("heap is full")
	}

	idx := h.len
	h.heap[idx] = id
	h.elems[idx] = elem
	h.len++
	for {
		parent := (idx - 1) / 2
		if parent == idx || h.heap[parent] >= h.heap[idx] {
			break
		}
		h.heap[parent], h.heap[idx] = h.heap[idx], h.heap[parent]
		h.elems[parent], h.elems[idx] = h.elems[idx], h.elems[parent]
		idx = parent
	}
}