Go
by smallnest

HeapMap, 一个混合功能的数据结构Go语言实现

今天在准备《秘而不宣》系列下一篇文章时，思绪飘散了，突然想到使用 Heap 的功能再加 HashTable (Map) 的功能，可以构造一种新的数据结构，然后把我聚合程序中的数据聚合数据结构替换掉，总之思绪翩翩。然后在网上搜了一下，这种数据结构其实早就有了，名字叫 HeapMap。

HeapMap (也叫做 PriorityMap) 是一种结合了堆和哈希映射的数据结构，常用于需要按键排序并进行高效查找的场景。它可以在优先级队列的基础上，使用哈希映射来提供快速访问和更新。HeapMap 在实现过程中利用堆的有序性和哈希表的快速查找能力，以支持按键排序和常数时间查找。

Go 语言支付 Rob Pike 在他的 Rob Pike's 5 Rules of Programming 第 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 的主要特点

  1. 堆的特点：HeapMap 内部通过堆来维护键的顺序，可以快速获取最小或最大键。堆提供了插入和删除堆顶元素的 O(log n) 时间复杂度。
  2. 哈希映射的特点：HeapMap 同时使用哈希映射以支持快速查找。哈希映射的查找、插入、删除等操作在理想情况下时间复杂度为 O(1)。
  3. 用途：HeapMap 适合需要频繁按键排序和快速查找的场景，比如带有优先级的缓存、调度系统、任务优先队列等。

HeapMap 的基本结构

  • 堆（Heap）：用来维持按键的顺序，堆可以是最小堆或最大堆，根据具体需求决定。
  • 哈希映射（Map）：用来存储每个键值对，并支持通过键快速查找元素。

你使用一个 container/heap + map 很容易实现一个 HeapMap, 其实我们没必要自己去写一个重复的轮子了，网上其他语言比如 Rust、Java 都有现成的实现，Go 语言中也有一个很好的实现：nemars/heapmap

HeapMap 的实现

nemars/heapmap 这个库是去年增加到 github 中的，我是第一个 star 它的人。我们看看它是怎么实现的。

结构体定义

1
2
3
4
5
6
7
8
9
10
11
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 方法
1
2
3
func (hm *heapmap[K, V, P]) Len() int {
	return len(hm.m)
}

读取h字段或者m字段的长度都可以。

Peek 方法

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

1
2
3
4
5
6
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元素。

1
2
3
4
5
6
7
8
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 方法来实现的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
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 方法就是获取任意的元素。

1
2
3
4
5
6
7
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 来保护它。