u

data-structure/ds-heap.go

@@ -0,0 +1,101 @@
+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
+	}
+}