HeapMap: 一個混合功能的數據結構 Go 語言實現

今天在準備《祕而不宣》系列下一篇文章時，思緒飄散了，突然想到使用 Heap 的功能再加 HashTable (Map) 的功能，可以構造一種新的數據結構，然後把我聚合程序中的數據聚合數據結構替換掉，總之思緒翩翩。然後在網上搜了一下，這種數據結構其實早就有了，名字叫 HeapMap。

HeapMap (也叫做 PriorityMap) 是一種結合了堆和哈希映射的數據結構，常用於需要按鍵排序並進行高效查找的場景。它可以在優先級隊列的基礎上，使用哈希映射來提供快速訪問和更新。HeapMap 在實現過程中利用堆的有序性和哈希表的快速查找能力，以支持按鍵排序和常數時間查找。

Go 語言支付 Rob Pike 在他的 Rob Pike's 5 Rules of Programming[1] 第 5 條就指出：

Data dominates. If you've chosen the right data structures and organized things well, the algorithms will almost always be self-evident. Data structures, not algorithms, are central to programming. 數據爲王。如果你選擇了合適的數據結構並進行了良好的組織，算法通常會變得顯而易見。編程的核心在於數據結構，而非算法。

所以，如果在合適的場景下，針對它的特點，使用 HeapMap 會取得事半功倍的效果。

`HeapMap` 的主要特點

堆的特點：HeapMap 內部通過堆來維護鍵的順序，可以快速獲取最小或最大鍵。堆提供了插入和刪除堆頂元素的 O(log n) 時間複雜度。
哈希映射的特點：HeapMap 同時使用哈希映射以支持快速查找。哈希映射的查找、插入、刪除等操作在理想情況下時間複雜度爲 O(1)。
用途：HeapMap 適合需要_頻繁按鍵排序和快速查找_的場景，比如帶有優先級的緩存、調度系統、任務優先隊列等。

`HeapMap` 的基本結構

堆（Heap）：用來維持按鍵的順序，堆可以是最小堆或最大堆，根據具體需求決定。
哈希映射（Map）：用來存儲每個鍵值對，並支持通過鍵快速查找元素。

你使用一個 container/heap + map 很容易實現一個 HeapMap, 其實我們沒必要自己去寫一個重複的輪子了，網上其他語言比如 Rust、Java 都有現成的實現，Go 語言中也有一個很好的實現：nemars/heapmap[2]

`HeapMap` 的實現

nemars/heapmap 這個庫是去年增加到 github 中的，我是第一個 star 它的人。我們看看它是怎麼實現的。

結構體定義

type Entry[K comparable, V, P any] struct {
 Key      K
 Value    V
 Priority P
 index    int
}

type heapmap[K comparable, V, P any] struct {
 h pq[K, V, P]
 m map[K]*Entry[K, V, P]
}

Entry 代表這個數據結構中的一個節點 (元素、條目) , 它包含 key、value 值，還有優先級，index 記錄它在堆的實現數組中的索引。

heapmap 代表 HeapMap 的實現，它包含兩個字段，第一個字段其實就是 Heap 的實現，爲了方便實現泛型，它就自己實現了一個堆。第二個字段就是一個 map 對象了。

典型的方法

數據結構定義清楚了，那就就可以實現它的方法了。它實現了一些便利的方法，我們值關注幾個實現就好了。

Len 方法

func (hm *heapmap[K, V, P]) Len() int {
 return len(hm.m)
}

讀取h字段或者m字段的長度都可以。

Peek 方法

返回 root 元素。最小堆就是返回最小的元素，最大堆就是返回最大的元素。

func (hm *heapmap[K, V, P]) Peek() (Entry[K, V, P], bool) {
 if hm.Empty() {
  return Entry[K, V, P]{}, false
 }
 return *hm.h.entries[0], true
}

Pop 方法

彈出 root 元素。

func (hm *heapmap[K, V, P]) Pop() (Entry[K, V, P], bool) {
 if hm.Empty() {
  return Entry[K, V, P]{}, false
 }
 e := *heap.Pop(&hm.h).(*Entry[K, V, P])
 delete(hm.m, e.Key)
 return e, true
}

注意涉及到元素的刪除操作，要同時刪除 map 中的元素。

Push 方法 (Set 方法)

其實作者沒有實現 Push 方法，而是使用 Set 方法來實現的。

func (hm *heapmap[K, V, P]) Set(key K, value V, priority P) {
 if e, ok := hm.m[key]; ok {
  e.Value = value
  e.Priority = priority
  heap.Fix(&hm.h, e.index)
  return
 }
 e := &Entry[K, V, P]{
  Key:      key,
  Value:    value,
  Priority: priority,
 }
 heap.Push(&hm.h, e)
 hm.m[key] = e
}

Set 方法有兩個功能。如果元素的 Key 已經存在，那麼就是更新元素，並且根據優先級進行調整。如果元素的 Key 不存在，那麼就是插入元素。

Get 方法

Get 方法就是獲取任意的元素。

func (hm *heapmap[K, V, P]) Get(key K) (Entry[K, V, P], bool) {
 if e, ok := hm.m[key]; ok {
  return *e, true
 }
 return Entry[K, V, P]{}, false
}

有一點你需要注意的是，這個數據結構不是線程安全的，如果你需要線程安全的話，你可以使用 sync.Mutex/sync.RWMutex 來保護它。

參考資料

[1]

Rob Pike's 5 Rules of Programming: https://users.ece.utexas.edu/~adnan/pike.html

[2]

nemars/heapmap: https://github.com/nemars/heapmap

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/9-DO7A2nQDPMBcYamdB77g

HeapMap 的主要特點

HeapMap 的基本結構

HeapMap 的實現