| 
 | 1 | +/*  | 
 | 2 | +MIT License  | 
 | 3 | +
  | 
 | 4 | +Copyright (c) 2021 Florimond Husquinet  | 
 | 5 | +
  | 
 | 6 | +Permission is hereby granted, free of charge, to any person obtaining a copy  | 
 | 7 | +of this software and associated documentation files (the "Software"), to deal  | 
 | 8 | +in the Software without restriction, including without limitation the rights  | 
 | 9 | +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell  | 
 | 10 | +copies of the Software, and to permit persons to whom the Software is  | 
 | 11 | +furnished to do so, subject to the following conditions:  | 
 | 12 | +
  | 
 | 13 | +The above copyright notice and this permission notice shall be included in all  | 
 | 14 | +copies or substantial portions of the Software.  | 
 | 15 | +
  | 
 | 16 | +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR  | 
 | 17 | +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,  | 
 | 18 | +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE  | 
 | 19 | +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER  | 
 | 20 | +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,  | 
 | 21 | +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE  | 
 | 22 | +SOFTWARE.  | 
 | 23 | +*/  | 
 | 24 | + | 
 | 25 | +/* A generic implementation of a binary heap */  | 
 | 26 | +package heap  | 
 | 27 | + | 
 | 28 | +import "sync"  | 
 | 29 | + | 
 | 30 | +type Heap[T any] struct {  | 
 | 31 | +	mu      sync.RWMutex  | 
 | 32 | +	data    []T  | 
 | 33 | +	compare func(T, T) int  | 
 | 34 | +}  | 
 | 35 | + | 
 | 36 | +// NewHeap constructs a heap using the provided comparator.  | 
 | 37 | +// The comparator should return -1 if a < b, 0 if a == b, and 1 if a > b.  | 
 | 38 | +// If compare orders values in ascending order, the heap behaves as a min-heap.  | 
 | 39 | +// To build a max-heap, invert the comparator (e.g., return -compare(a, b)).  | 
 | 40 | +func NewHeap[T any](compare func(T, T) int) *Heap[T] {  | 
 | 41 | +	return &Heap[T]{  | 
 | 42 | +		data:    make([]T, 0),  | 
 | 43 | +		compare: compare,  | 
 | 44 | +	}  | 
 | 45 | +}  | 
 | 46 | + | 
 | 47 | +// NewHeapFromSlice builds a heap in O(n) from an initial slice.  | 
 | 48 | +func NewHeapFromSlice[T any](values []T, compare func(T, T) int) *Heap[T] {  | 
 | 49 | +	h := &Heap[T]{  | 
 | 50 | +		data:    append([]T(nil), values...),  | 
 | 51 | +		compare: compare,  | 
 | 52 | +	}  | 
 | 53 | +	// heapify bottom-up  | 
 | 54 | +	for i := (len(h.data) / 2) - 1; i >= 0; i-- {  | 
 | 55 | +		h.sinkDown(i)  | 
 | 56 | +	}  | 
 | 57 | +	return h  | 
 | 58 | +}  | 
 | 59 | + | 
 | 60 | +// Peek returns the top element without removing it.  | 
 | 61 | +func (h *Heap[T]) Peek() (value T, ok bool) {  | 
 | 62 | +	h.mu.RLock()  | 
 | 63 | +	defer h.mu.RUnlock()  | 
 | 64 | +	if len(h.data) == 0 {  | 
 | 65 | +		return value, false  | 
 | 66 | +	}  | 
 | 67 | +	return h.data[0], true  | 
 | 68 | +}  | 
 | 69 | + | 
 | 70 | +func (h *Heap[T]) Len() int {  | 
 | 71 | +	h.mu.RLock()  | 
 | 72 | +	defer h.mu.RUnlock()  | 
 | 73 | +	return len(h.data)  | 
 | 74 | +}  | 
 | 75 | + | 
 | 76 | +func (h *Heap[T]) Push(value T) {  | 
 | 77 | +	h.mu.Lock()  | 
 | 78 | +	defer h.mu.Unlock()  | 
 | 79 | +	h.data = append(h.data, value)  | 
 | 80 | +	idx := len(h.data) - 1  | 
 | 81 | +	h.bubbleUp(idx)  | 
 | 82 | +}  | 
 | 83 | + | 
 | 84 | +func (h *Heap[T]) Pop() (value T, ok bool) {  | 
 | 85 | +	h.mu.Lock()  | 
 | 86 | +	defer h.mu.Unlock()  | 
 | 87 | +	n := len(h.data)  | 
 | 88 | +	if n == 0 {  | 
 | 89 | +		return value, false  | 
 | 90 | +	}  | 
 | 91 | +	top := h.data[0]  | 
 | 92 | +	h.data[0] = h.data[n-1]  | 
 | 93 | +	h.data = h.data[:n-1]  | 
 | 94 | +	h.sinkDown(0)  | 
 | 95 | +	return top, true  | 
 | 96 | +}  | 
 | 97 | + | 
 | 98 | +// Min heap: if a node is less than its parent, swap them.  | 
 | 99 | +func (h *Heap[T]) bubbleUp(index int) {  | 
 | 100 | +	if index == 0 {  | 
 | 101 | +		return  | 
 | 102 | +	}  | 
 | 103 | +	var parent = (index - 1) / 2  | 
 | 104 | +	if h.compare(h.data[index], h.data[parent]) < 0 {  | 
 | 105 | +		h.swap(index, parent)  | 
 | 106 | +		h.bubbleUp(parent)  | 
 | 107 | +	}  | 
 | 108 | +}  | 
 | 109 | + | 
 | 110 | +// Min heap: if a node is greater than its children, swap the node with the smallest child.  | 
 | 111 | +func (h *Heap[T]) sinkDown(index int) {  | 
 | 112 | +	n := len(h.data)  | 
 | 113 | +	left := index*2 + 1  | 
 | 114 | +	right := index*2 + 2  | 
 | 115 | +	smallest := index  | 
 | 116 | +	if left < n && h.compare(h.data[left], h.data[smallest]) < 0 {  | 
 | 117 | +		smallest = left  | 
 | 118 | +	}  | 
 | 119 | +	if right < n && h.compare(h.data[right], h.data[smallest]) < 0 {  | 
 | 120 | +		smallest = right  | 
 | 121 | +	}  | 
 | 122 | +	if smallest != index {  | 
 | 123 | +		h.swap(index, smallest)  | 
 | 124 | +		h.sinkDown(smallest)  | 
 | 125 | +	}  | 
 | 126 | +}  | 
 | 127 | + | 
 | 128 | +func (h *Heap[T]) swap(i, j int) {  | 
 | 129 | +	h.data[i], h.data[j] = h.data[j], h.data[i]  | 
 | 130 | +}  | 
0 commit comments