// 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 } }