102 lines
2.2 KiB
Go
102 lines
2.2 KiB
Go
package ds
|
|
|
|
// ds-heap.go
|
|
// 泛型二叉堆,手动实现 siftUp/siftDown,不依赖 container/heap。
|
|
// 通过 less 函数决定堆性质:less(a,b)=true 表示 a 应在 b 上方。
|
|
// - 传 a < b → 小顶堆
|
|
// - 传 a > b → 大顶堆
|
|
// Push/Pop 均 O(log n),Peek O(1)。
|
|
|
|
// Heap 泛型堆。零值不可直接用,必须用 NewHeap 创建。
|
|
type Heap[T any] struct {
|
|
data []T
|
|
less func(a, b T) bool
|
|
}
|
|
|
|
// NewHeap 创建堆。less(a,b) 返回 true 表示 a 优先级更高(更靠堆顶)。
|
|
func NewHeap[T any](less func(a, b T) bool) *Heap[T] {
|
|
return &Heap[T]{less: less}
|
|
}
|
|
|
|
// Len 元素个数。
|
|
func (h *Heap[T]) Len() int { return len(h.data) }
|
|
|
|
// Empty 是否为空。
|
|
func (h *Heap[T]) Empty() bool { return len(h.data) == 0 }
|
|
|
|
// Peek 查看堆顶,不移除。
|
|
func (h *Heap[T]) Peek() (T, bool) {
|
|
var zero T
|
|
if len(h.data) == 0 {
|
|
return zero, false
|
|
}
|
|
return h.data[0], true
|
|
}
|
|
|
|
// Push 插入元素:先追加到末尾,再 siftUp。
|
|
func (h *Heap[T]) Push(v T) {
|
|
h.data = append(h.data, v)
|
|
h.siftUp(len(h.data) - 1)
|
|
}
|
|
|
|
// Pop 弹出堆顶:把末尾元素换到堆顶,缩容,再 siftDown。
|
|
func (h *Heap[T]) Pop() (T, bool) {
|
|
var zero T
|
|
n := len(h.data)
|
|
if n == 0 {
|
|
return zero, false
|
|
}
|
|
top := h.data[0]
|
|
if n == 1 {
|
|
h.data = h.data[:0]
|
|
return top, true
|
|
}
|
|
h.data[0] = h.data[n-1]
|
|
h.data[n-1] = zero // 置零释放引用
|
|
h.data = h.data[:n-1]
|
|
h.siftDown(0)
|
|
return top, true
|
|
}
|
|
|
|
// Init 从已有切片原地建堆,O(n)。调用后 h 持有该切片的引用。
|
|
func (h *Heap[T]) Init(data []T) {
|
|
h.data = data
|
|
// 从最后一个非叶子节点开始向下调整。
|
|
for i := len(h.data)/2 - 1; i >= 0; i-- {
|
|
h.siftDown(i)
|
|
}
|
|
}
|
|
|
|
// siftUp 从 i 向上调整。
|
|
func (h *Heap[T]) siftUp(i int) {
|
|
for i > 0 {
|
|
parent := (i - 1) / 2
|
|
if !h.less(h.data[i], h.data[parent]) {
|
|
break
|
|
}
|
|
h.data[i], h.data[parent] = h.data[parent], h.data[i]
|
|
i = parent
|
|
}
|
|
}
|
|
|
|
// siftDown 从 i 向下调整。
|
|
func (h *Heap[T]) siftDown(i int) {
|
|
n := len(h.data)
|
|
for {
|
|
l := 2*i + 1
|
|
r := 2*i + 2
|
|
best := i
|
|
if l < n && h.less(h.data[l], h.data[best]) {
|
|
best = l
|
|
}
|
|
if r < n && h.less(h.data[r], h.data[best]) {
|
|
best = r
|
|
}
|
|
if best == i {
|
|
break
|
|
}
|
|
h.data[i], h.data[best] = h.data[best], h.data[i]
|
|
i = best
|
|
}
|
|
}
|