// Package binheap implements a 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 } }