116 lines
2.7 KiB
Go
116 lines
2.7 KiB
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
|
|
}
|
|
}
|