103 lines
2.2 KiB
Go
103 lines
2.2 KiB
Go
package binheap
|
|
|
|
import "golang.org/x/exp/constraints"
|
|
|
|
// H is a generic, non-concurrent 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
|
|
}
|
|
}
|