priorityq/binheap/lib.go

// Package binheap implements a non-concurrent binary max-heap.
//
// # Implementation
//
// [H] is parameterized over two types, one for the priority levels, one for
// the elements. Internally, there are two equally-sized buffers for these
// types. Re-heaping operations swap corresponding entries in these buffers
// in lock-step.
package binheap

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

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

// Make creates a new heap.
func Make[P constraints.Ordered, E any](cap int) H[P, E] {
	priorities := make([]P, cap)
	elements := make([]E, cap)
	h := H[P, E]{prs: priorities, els: elements}
	return h
}

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

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

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

// CanInsert returns true if the heap has unused capacity, otherwise false.
func (h *H[P, E]) CanInsert() bool {
	return cap(h.prs)-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[P, E]) Extract() (P, E) {
	if !h.CanExtract() {
		panic("heap is empty")
	}
	// extract root
	priority := h.prs[0]
	element := h.els[0]
	// move last entry to root position
	h.prs[0] = h.prs[h.len-1]
	h.els[0] = h.els[h.len-1]
	// clear the former last entry position,
	// so as not to hold onto garbage
	var emptyPriority P
	var emptyElem E
	h.prs[h.len-1] = emptyPriority
	h.els[h.len-1] = emptyElem
	// heap-down
	h.len--
	idx := 0
	for {
		left := idx<<1 + 1
		right := idx<<1 + 2
		largest := idx
		if left < h.len && h.prs[left] > h.prs[largest] {
			largest = left
		}
		if right < h.len && h.prs[right] > h.prs[largest] {
			largest = right
		}
		if largest == idx {
			break
		}
		h.prs[idx], h.prs[largest] = h.prs[largest], h.prs[idx]
		h.els[idx], h.els[largest] = h.els[largest], h.els[idx]
		idx = largest
	}

	return priority, element
}

// Insert adds an item to the heap, then performs a heap-up pass.
//
// If the heap is full, it panics.
func (h *H[P, E]) Insert(priority P, elem E) {
	if !h.CanInsert() {
		panic("heap is full")
	}
	// insert new item into last position
	idx := h.len
	h.prs[idx] = priority
	h.els[idx] = elem
	// heap-up
	h.len++
	for {
		parent := (idx - 1) >> 1
		if parent < 0 || h.prs[parent] >= h.prs[idx] {
			break
		}
		h.prs[parent], h.prs[idx] = h.prs[idx], h.prs[parent]
		h.els[parent], h.els[idx] = h.els[idx], h.els[parent]
		idx = parent
	}
}