理解 Golang 中的最大/最小堆、`heap` 與優先隊列
最大堆、最小堆、 heap 、 優先隊列在數據結構算法題目里都是一個東西。這里討論 container/heap 的使用。
參考:
- https://pkg.go.dev/container/heap
- https://github.com/EndlessCheng/codeforces-go/blob/master/copypasta/heap.go 靈佬筆記,非常有用
在算法題目中,我們經常遇到需要動態維護一個集合的最值(最大或最小值)的場景,例如找出動態數據流中的第 K 大元素、Dijkstra 算法中的節點松弛操作等。這些場景的共同特點是,我們不僅需要找到當前的最值,還需要高效地添加新元素和刪除最值。 優先隊列 (Priority Queue) 是實現這種操作的理想抽象數據結構,而 堆 (heap) 則是實現優先隊列最常用、最高效的具體數據結構。
Golang 的標準庫 container/heap 提供了一套堆操作的算法。需要注意的是,它并沒有提供一個可以直接使用的、開箱即用的堆類型,而是定義了一個需要用戶自己實現的接口 heap.Interface 。用戶需要提供一個滿足該接口的數據類型(通常是一個切片),container/heap 包中的函數,如 heap.Push 和 heap.Pop ,會基于用戶提供的類型來維護堆的性質。
這種設計體現了 Go 語言接口的哲學:定義行為,而不是具體實現。它給予了開發者極大的靈活性,讓我們可以對任意類型的集合實現堆操作,無論是整數、字符串還是自定義的結構體。
heap.Interface 與官方示例
要使用 container/heap 包,我們首先需要理解它所依賴的核心接口——heap.Interface。
// A Interface is a type that can be used as a min-heap.
// Methods of this interface are documented in the heap package.
type Interface interface {
sort.Interface // 內嵌了 sort.Interface
Push(x any) // add x as element Len()
Pop() any // remove and return element Len() - 1.
}
可以看到,heap.Interface 內嵌了 sort.Interface。這意味著任何想要實現堆操作的類型,都必須首先實現 sort.Interface,即以下三個方法:
- Len() int: 返回集合中元素的數量。
- Less(i, j int) bool: 比較索引 i 和 j 處的元素。如果 h[i] 應該排在 h[j] 前面,則返回 true。對于 最小堆 ,這意味著 h[i] < h[j];對于 最大堆 ,則是 h[i] > h[j]。
- Swap(i, j int): 交換索引 i 和 j 處的元素。
除此之外,還需要實現兩個特定于堆的方法:
- Push(x any): 在集合的“末尾”添加一個新元素 x。通常,這意味著將元素 append 到切片的末尾。
- Pop() any: 從集合的“末尾”移除并返回一個元素。通常,這意味著移除并返回切片的最后一個元素。
一個常見的困惑 :為什么 Pop 方法是移除 最后一個 元素,而不是堆頂元素(索引為 0 的元素)?這正是 container/heap 包算法設計的巧妙之處。當我們調用 heap.Pop(h) 時,該函數會:
- 首先將堆頂元素(h[0])與堆的最后一個元素(h[len(h)-1])交換位置。
- 此時,原本的最值(最小或最大元素)已經被移動到了切片的末尾。
- 然后,heap.Pop 會調用我們自己實現的 Pop() 方法。我們的 Pop() 方法只需要簡單地移除并返回切片的最后一個元素即可,這個元素正是我們所需要的原堆頂元素。
- 最后,heap.Pop 內部會調整堆,使得除最后一個元素外,剩下的 n-1 個元素重新滿足堆的性質。
接下來,我們通過官方文檔的兩個例子來具體理解這個過程。
示例一:整數最小堆
這個例子展示了如何基于 []int 切片構建一個整數最小堆。
// 該示例演示了如何使用 heap 接口構建一個整數最小堆。
package main
import (
"container/heap"
"fmt"
)
// IntHeap 是一個整數最小堆。它本質上是一個 int 類型的切片。
type IntHeap []int
// Len 返回切片的長度。
func (h IntHeap) Len() int { return len(h) }
// Less 用于定義元素間的大小關系。對于最小堆,如果 h[i] < h[j],則返回 true。
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] }
// Swap 交換切片中兩個元素的位置。
func (h IntHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
// Push 在切片末尾添加一個元素。
// 注意:Push 和 Pop 方法使用指針接收者 (*IntHeap),因為它們需要修改切片的長度,而不僅僅是內容。
func (h *IntHeap) Push(x any) {
*h = append(*h, x.(int))
}
// Pop 從切片末尾移除并返回元素。
func (h *IntHeap) Pop() any {
old := *h
n := len(old)
x := old[n-1]
*h = old[0 : n-1]
return x
}
// 這個例子向 IntHeap 中插入了幾個整數,檢查了最小值,并按優先級順序將它們移除。
func main() {
// 創建一個 IntHeap 實例,并初始化。
h := &IntHeap{2, 1, 5}
heap.Init(h) // 將無序的切片整理成一個最小堆。
// 向堆中推入一個新元素。
heap.Push(h, 3)
// 堆頂元素 h[0] 即為最小值。
fmt.Printf("minimum: %d\n", (*h)[0]) // 輸出 "minimum: 1"
// 持續從堆中彈出元素,直到堆為空。
// 彈出的順序將是從小到大。
for h.Len() > 0 {
fmt.Printf("%d ", heap.Pop(h)) // 輸出 "1 2 3 5 "
}
}
minimum: 1
1 2 3 5
示例二:使用堆實現優先隊列
這個例子更進一步,展示了如何用堆實現一個管理復雜對象的優先隊列,并且支持在隊列中修改元素的優先級。
// 該示例演示了如何使用 heap 接口構建一個優先隊列。
package main
import (
"container/heap"
"fmt"
)
// Item 是我們在優先隊列中管理的元素。
type Item struct {
value string // 元素的值,可以是任意類型。
priority int // 元素在隊列中的優先級。
// index 字段對于 update 方法至關重要。
// 它由 heap.Interface 的方法(特別是 Swap)來維護。
index int // 元素在堆中的索引。
}
// PriorityQueue 實現了 heap.Interface 接口,并持有 Item 類型的元素。
// 它是一個 Item 指針的切片。
type PriorityQueue []*Item
func (pq PriorityQueue) Len() int { return len(pq) }
func (pq PriorityQueue) Less(i, j int) bool {
// 我們希望 Pop 返回的是最高優先級的元素,而不是最低的,
// 所以這里我們使用 > (大于號)。
return pq[i].priority > pq[j].priority
}
func (pq PriorityQueue) Swap(i, j int) {
// 交換元素
pq[i], pq[j] = pq[j], pq[i]
// **關鍵**:交換元素后,必須同時更新它們在堆中的 index。
pq[i].index = i
pq[j].index = j
}
func (pq *PriorityQueue) Push(x any) {
n := len(*pq)
item := x.(*Item)
// 設置新元素的初始 index。
item.index = n
*pq = append(*pq, item)
}
func (pq *PriorityQueue) Pop() any {
old := *pq
n := len(old)
item := old[n-1]
old[n-1] = nil // 避免內存泄漏
item.index = -1 // for safety
*pq = old[0 : n-1]
return item
}
// update 方法修改隊列中一個 Item 的優先級和值。
func (pq *PriorityQueue) update(item *Item, value string, priority int) {
item.value = value
item.priority = priority
// 在修改了優先級后,元素在堆中的位置可能不再正確,
// 因此需要調用 heap.Fix 來恢復堆的屬性。
heap.Fix(pq, item.index)
}
func main() {
// 一些元素和它們的優先級。
items := map[string]int{
"banana": 3, "apple": 2, "pear": 4,
}
// 創建優先隊列,將元素放入其中,并建立堆的不變性。
pq := make(PriorityQueue, len(items))
i := 0
for value, priority := range items {
pq[i] = &Item{
value: value,
priority: priority,
index: i,
}
i++
}
heap.Init(&pq) // 初始化堆
// 插入一個新元素,然后修改它的優先級。
item := &Item{
value: "orange",
priority: 1,
}
heap.Push(&pq, item)
pq.update(item, item.value, 5) // 將 orange 的優先級從 1 更新到 5
// 按優先級從高到低的順序取出所有元素。
for pq.Len() > 0 {
item := heap.Pop(&pq).(*Item)
fmt.Printf("%.2d:%s ", item.priority, item.value)
// 輸出: "05:orange 04:pear 03:banana 02:apple "
}
}
05:orange 04:pear 03:banana 02:apple
container/heap 核心 API
現在我們來詳細解釋一下 heap 包提供的幾個核心函數:
- func Init(h Interface) 此函數用于將一個無序的、滿足 Interface 的數據集合整理成一個合法的堆。它的實現方式是從最后一個非葉子節點開始,自下而上、自右向左地對每個子樹調用 down(一個內部函數)進行調整,使其滿足堆的性質。時間復雜度為 O(n),比逐個 Push 元素(O(n log n))更高效。
- func Push(h Interface, x any) 向堆 h 中添加一個新元素 x。它首先調用用戶定義的 Push(x) 方法將元素添加到數據集合的末尾,然后調用 up(一個內部函數)將這個新元素“上浮”到它在堆中的正確位置。時間復雜度為 O(log n)。
- func Pop(h Interface) any 從堆 h 中移除并返回堆頂元素(最小值或最大值)。如前所述,它通過將堆頂元素與最后一個元素交換,然后調用用戶定義的 Pop() 方法來實現。之后,它會調用 down 將新的堆頂元素“下沉”到正確位置,以維持堆的性質。時間復雜度為 O(log n)。
- func Remove(h Interface, i int) any 移除并返回堆中索引為 i 的元素。這是一個更通用的刪除操作。它的實現方式是將索引 i 的元素與最后一個元素交換,然后移除最后一個元素(即我們想刪除的元素),最后對交換到位置 i 的新元素進行調整(可能上浮也可能下沉)來恢復堆的性質。時間復雜度為 O(log n)。
- func Fix(h Interface, i int) 當用戶直接修改了堆中索引為 i 的元素的值(比如 PriorityQueue 例子中的 update 操作),這個元素的位置可能就不再滿足堆的性質了。此時應調用 Fix(h, i),該函數會自動地將該元素上浮或下沉到它新的正確位置,從而恢復整個堆的性質。i 就是被修改元素在底層切片中的索引。
堆的內存布局
我們有必要先理解堆在內存中是如何存儲的。
從邏輯上講,堆是一個 完全二叉樹 (Complete Binary Tree) 。這意味著除了最底層外,其他層都是完全填滿的,并且最底層的節點都盡可能地靠左排列。然而,在物理存儲上,堆并不會像鏈表那樣使用指針來連接父子節點。相反,它被巧妙地存儲在一個連續的內存空間中,比如 Golang 中的 切片 (slice) 或數組。
這種設計帶來了巨大的性能優勢:它不需要額外的指針開銷,并且由于數據是連續存儲的,可以更好地利用 CPU 緩存,提高訪問效率。
我們通過切片的索引 i 就可以計算出任意節點的父節點和子節點的索引:
- 假設一個節點的索引為 i(i 從 0 開始):
- 它的左子節點的索引是 2*i + 1
- 它的右子節點的索引是 2*i + 2
- 它的父節點的索引是 (i - 1) / 2 (整數除法,自動向下取整)
例如,對于一個最小堆,其切片為 [3, 5, 8, 10, 7],它的邏輯樹形結構如下:
邏輯樹形結構 物理切片存儲
3 (i=0) Index: 0 1 2 3 4
/ \ Value: [3, 5, 8, 10, 7]
/ \
5 (i=1) 8 (i=2)
/ \
/ \
10(i=3) 7(i=4)
- 節點 3 (i=0) 的左子節點是 2*0 + 1 = 1,即 5;右子節點是 2*0 + 2 = 2,即 8。
- 節點 5 (i=1) 的父節點是 (1 - 1) / 2 = 0,即 3。
container/heap 包中的所有算法,如 up 和 down,都是基于這個索引計算規則來操作底層切片,從而高效地維護堆的邏輯結構和性質。
實用模板與技巧
在解決算法問題時,為了提高編碼效率,我們可以定義一些可復用的堆模板。
模板一:利用內嵌 sort.IntSlice
sort.IntSlice 已經為 []int 實現了 Len、Less 和 Swap 方法。通過在我們的堆類型中 內嵌 (embedding)sort.IntSlice,我們可以自動獲得這些方法的實現,從而只需要編寫 Push 和 Pop 即可。
內嵌語法解釋 :在 struct 中寫下一個沒有字段名的類型(如 type hp struct{ sort.IntSlice }),就是內嵌。這使得 hp 的實例可以直接訪問 sort.IntSlice 的方法和字段。在 Push 方法中,h.IntSlice = append(h.IntSlice, v.(int)) 這行代碼中,h.IntSlice 就是訪問內嵌的 sort.IntSlice 字段,它本身就是一個 []int。
package main
import (
"container/heap"
"fmt"
"sort"
)
// hp 是一個最小堆。
// 它內嵌了 sort.IntSlice,自動獲得了 Len, Less, Swap 方法。
type hp struct{ sort.IntSlice }
// 如果需要最大堆,只需覆蓋 Less 方法即可。
// func (h hp) Less(i, j int) bool { return h.IntSlice[i] > h.IntSlice[j] }
func (h *hp) Push(v any) {
h.IntSlice = append(h.IntSlice, v.(int))
}
func (h *hp) Pop() any {
a := h.IntSlice
v := a[len(a)-1]
h.IntSlice = a[:len(a)-1]
return v
}
// 為了方便使用,可以封裝類型安全的 push 和 pop 方法。
func (h *hp) push(v int) {
heap.Push(h, v)
}
func (h *hp) pop() int {
return heap.Pop(h).(int)
}
// replace 彈出并返回堆頂,同時將新元素 v 入堆。
// 相比 pop + push,效率更高,因為它只需要一次堆調整。
// 要求 h 非空。
func (h *hp) replace(v int) int {
top := h.IntSlice[0]
h.IntSlice[0] = v
heap.Fix(h, 0) // 調整堆頂元素
return top
}
// pushPop 先將 v 入堆,然后彈出并返回堆頂。
func (h *hp) pushPop(v int) int {
// 如果新元素 v 比堆頂還小(最小堆),
// 那么 v 將成為新的堆頂并被立即彈出,堆本身不變。
// 所以只有當 v > h.IntSlice[0] 時才需要操作。
if h.Len() > 0 && v > h.IntSlice[0] { // 最大堆則為 v < h.IntSlice[0]
v, h.IntSlice[0] = h.IntSlice[0], v // 交換新元素和堆頂
heap.Fix(h, 0) // 調整新的堆頂
}
return v
}
func main() {
// 這是一個可以直接運行的 hp 示例
minHeap := &hp{}
minHeap.push(4)
minHeap.push(1)
minHeap.push(7)
fmt.Println("堆頂(最小值):", (*minHeap).IntSlice[0]) // 輸出 1
popped := minHeap.pop()
fmt.Println("彈出:", popped) // 輸出 1
fmt.Println("新的堆頂:", (*minHeap).IntSlice[0]) // 輸出 4
replacedVal := minHeap.replace(0)
fmt.Println("被替換的堆頂:", replacedVal) // 輸出 4
fmt.Println("替換后的堆頂:", (*minHeap).IntSlice[0]) // 輸出 0
}
堆頂(最小值): 1
彈出: 1
新的堆頂: 4
被替換的堆頂: 4
替換后的堆頂: 0
模板二:自定義類型堆
當我們需要處理 int32、float64 或其他非 int 類型時,只需定義一個新的類型并實現完整的 heap.Interface。
package main
import (
"container/heap"
"fmt"
)
// 自定義 int32 類型的最小堆
type hp32 []int32
func (h hp32) Len() int { return len(h) }
func (h hp32) Less(i, j int) bool { return h[i] < h[j] } // > 為最大堆
func (h hp32) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *hp32) Push(v any) { *h = append(*h, v.(int32)) }
func (h *hp32) Pop() any {
a := *h
v := a[len(a)-1]
*h = a[:len(a)-1]
return v
}
func (h *hp32) push(v int32) { heap.Push(h, v) }
func (h *hp32) pop() int32 { return heap.Pop(h).(int32) }
func main() {
// 這是一個可以直接運行的 hp32 示例
h := &hp32{100, 50, 200}
heap.Init(h)
fmt.Println("初始化后的堆頂:", (*h)[0]) // 輸出 50
h.push(25)
fmt.Println("Push 25 后的堆頂:", (*h)[0]) // 輸出 25
fmt.Println("按順序彈出:")
for h.Len() > 0 {
fmt.Printf("%d ", h.pop()) // 輸出 25 50 100 200
}
fmt.Println()
}
初始化后的堆頂: 50
Push 25 后的堆頂: 25
按順序彈出:
25 50 100 200
模板三:支持修改與刪除的堆
這個模板借鑒了官方 PriorityQueue 的思路,通過在堆中存儲指針,并維護每個元素在堆中的索引,從而實現了對堆中任意元素的修改和刪除。這在一些復雜的算法題中非常有用。
package main
import (
"container/heap"
"fmt"
)
// viPair 包含一個值 v 和它在堆中的索引 hi。
type viPair struct {
v int
hi int // *viPair 在 mh 中的下標,可隨著 Push/Pop 等操作自動改變
}
// mh (modifiable heap) 是一個支持修改的最小堆。
type mh []*viPair
func (h mh) Len() int { return len(h) }
func (h mh) Less(i, j int) bool { return h[i].v < h[j].v } // > 為最大堆
func (h mh) Swap(i, j int) {
h[i], h[j] = h[j], h[i]
// 關鍵:交換元素后,必須更新它們的索引 hi
h[i].hi = i
h[j].hi = j
fmt.Println(h[i], h[j])
}
func (h *mh) Push(v any) { *h = append(*h, v.(*viPair)) }
func (h *mh) Pop() any {
a := *h
v := a[len(a)-1]
*h = a[:len(a)-1]
return v
}
// push 會返回一個指向新元素的指針,外部可以持有該指針。
func (h *mh) push(v int) *viPair {
p := &viPair{v: v, hi: len(*h)}
heap.Push(h, p)
return p
}
func (h *mh) pop() *viPair { return heap.Pop(h).(*viPair) }
func (h *mh) fix(i int) { heap.Fix(h, i) }
func (h *mh) remove(i int) *viPair { return heap.Remove(h, i).(*viPair) }
func main() {
// 這是一個可以直接運行的 mh 示例
h := &mh{}
// 推入元素并保存它們的指針(句柄)
p1 := h.push(10)
h.push(5)
p3 := h.push(15)
fmt.Printf("初始堆頂: %d (索引 %d)\n", (*h)[0].v, (*h)[0].hi) // 5 (索引 0)
// 修改 p1 的值,它當前不在堆頂
fmt.Printf("修改前 p1 的值: %d, 在堆中索引為: %d\n", p1.v, p1.hi)
p1.v = 2 // 將 p1 的值從 10 改為 2
h.fix(p1.hi) // 修復堆
fmt.Printf("修改后 p1 的值: %d, 在堆中索引為: %d\n", p1.v, p1.hi)
fmt.Printf("修改后的堆頂: %d (索引 %d)\n", (*h)[0].v, (*h)[0].hi) // 2 (索引 0)
// 刪除 p3
fmt.Printf("刪除前 p3 的值: %d, 在堆中索引為: %d\n", p3.v, p3.hi)
removed := h.remove(p3.hi)
fmt.Printf("被刪除的元素值: %d\n", removed.v)
fmt.Println("按順序彈出剩余元素:")
for h.Len() > 0 {
p := h.pop()
fmt.Printf("值: %d, 彈出前索引: %d\n", p.v, p.hi)
}
}
初始堆頂: 5 (索引 0)
修改前 p1 的值: 10, 在堆中索引為: 1
修改后 p1 的值: 2, 在堆中索引為: 0
修改后的堆頂: 2 (索引 0)
刪除前 p3 的值: 15, 在堆中索引為: 2
被刪除的元素值: 15
按順序彈出剩余元素:
值: 2, 彈出前索引: 1
值: 5, 彈出前索引: 0
為什么值為 2 的元素在作為堆頂被彈出時,其 hi(彈出前索引)字段顯示為 1 而不是 0?
這個現象確實看起來有悖直覺,但它恰恰揭示了 heap.Pop 操作和我們自定義 Swap 方法聯動的精確過程。讓我們一步步拆解 h.pop() 調用時發生了什么:
當時堆的切片狀態是:[{v:2, hi:0}, {v:5, hi:1}]。
- 我們調用了 h.pop(),它內部調用了 heap.Pop(h)。
- heap.Pop(h) 的首要任務是把堆頂元素(我們想要的返回值)取出來。它的策略是:將堆頂元素 h[0] 與堆的最后一個元素 h[len-1](在這里是 h[1])進行交換。
- 這個交換操作觸發了我們定義的 Swap(0, 1) 方法。
- 在 Swap(0, 1) 方法內部:
- 現在位于索引 0 的元素(值是 5)的 hi 被更新為 0。
- 現在位于索引 1 的元素(值是 2)的 hi 被更新為 1。
- h[0] 和 h[1] 的指針被交換。交換后,切片在內存中的狀態變為:[{v:5, hi:1}, {v:2, hi:0}]。
- 關鍵來了 :緊接著,Swap 方法更新了這兩個被交換元素的 hi 字段以反映它們在切片中的 新位置。
- Swap 執行完畢。此時,我們即將彈出的、值為 2 的元素,它正位于切片的末尾(索引 1),并且它自身的 hi 字段剛剛被更新為了 1。
- heap.Pop 接著調用我們定義的 Pop() 方法,該方法從切片末尾移除并返回 h[1],即 {v:2, hi:1} 這個 viPair 實例。
- 因此,fmt.Printf 打印出的 p.v 是 2,p.hi 是 1。
結論 :這個輸出是完全正確的。hi 字段忠實地記錄了元素在被 Pop 方法從切片中正式移除前的最后一刻,它在切片中的索引位置。這個位置因為 heap.Pop 內部的交換操作而從 0 變成了 1。這也凸顯了 hi 字段的動態性——它總是在 Swap 操作后被立即更新,以保證其值的實時準確性。