前 K 个高频元素
本题主要使用 Go 的 heap.Interface 接口来实现一个小顶堆来并解答 TopK 问题,问题链接:
347. 前 K 个高频元素 – 力扣(LeetCode) (leetcode-cn.com)
包 heap 为所有实现了 heap.Interface 的类型提供堆操作


因此要实现 heap.Interface 接口,需要实现上面 5 个方法
—- 我是分割线 —-
一个堆即是一棵树, 这棵树的每个节点的值都比它的子节点的值要小( 小顶堆 ), 而整棵树最小的值位于树根(root), 也即是索引 0 的位置上。
堆是实现优先队列的一种常见方法。 为了构建优先队列, 用户在实现堆接口时, 需要让 Less() 方法返回逆序的结果, 这样就可以在使用 Push 添加元素的同时, 通过 Pop 移除队列中优先级最高的元素了。 具体的实现请看接下来展示的优先队列例子。
package stackqueue
import (
"container/heap"
)
// Item 是优先队列中包含的元素。
type Item struct {
value int // 元素的值,可以是任意字符串。
priority int // 元素在队列中的优先级。
// 元素的索引可以用于更新操作,它由 heap.Interface 定义的方法维护。
index int // 元素在堆中的索引。
}
// PriorityQueue 一个实现了 heap.Interface 接口的优先队列,队列中包含任意多个 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]
pq[i].index = i
pq[j].index = j
}
func (pq *PriorityQueue) Push(x interface{}) {
n := len(*pq)
item := x.(*Item)
item.index = n
*pq = append(*pq, item)
}
func (pq *PriorityQueue) Pop() interface{} {
old := *pq
n := len(old)
item := old[n-1]
item.index = -1 // 为了安全性考虑而做的设置
*pq = old[0 : n-1]
return item
}
// 更新函数会修改队列中指定元素的优先级以及值。
func (pq *PriorityQueue) update(item *Item, value int, priority int) {
item.value = value
item.priority = priority
heap.Fix(pq, item.index)
}
func topKFrequent(nums []int, k int) []int {
frequency := make(map[int]int)
for _, n := range nums {
frequency[n]++
}
pq := make(PriorityQueue, 0)
heap.Init(&pq)
for key, val := range frequency {
heap.Push(&pq, &Item{value: key, priority: val})
if pq.Len() > k {
heap.Pop(&pq)
}
}
res := make([]int, 0)
for pq.Len() > 0 {
res = append(res, heap.Pop(&pq).(*Item).value)
}
return res
}
Go 中实现接口时,Receiver 可以是类型本身(虽说slice本身类似指针),也可以是类型的指针
下面这种写法亦可(update 不是 heap.Interface 需要实现):
package stackqueue
import (
"container/heap"
)
// Item 是优先队列中包含的元素。
type Item struct {
value int // 元素的值,可以是任意字符串。
priority int // 元素在队列中的优先级。
// 元素的索引可以用于更新操作,它由 heap.Interface 定义的方法维护。
index int // 元素在堆中的索引。
}
// PriorityQueue 一个实现了 heap.Interface 接口的优先队列,队列中包含任意多个 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]
(*pq)[i].index = i
(*pq)[j].index = j
}
func (pq *PriorityQueue) Push(x interface{}) {
n := len(*pq)
item := x.(*Item)
item.index = n
*pq = append(*pq, item)
}
func (pq *PriorityQueue) Pop() interface{} {
old := *pq
n := len(old)
item := old[n-1]
item.index = -1 // 为了安全性考虑而做的设置
*pq = old[0 : n-1]
return item
}
func topKFrequent(nums []int, k int) []int {
frequency := make(map[int]int)
for _, n := range nums {
frequency[n]++
}
pq := make(PriorityQueue, 0)
heap.Init(&pq)
for key, val := range frequency {
heap.Push(&pq, &Item{value: key, priority: val})
if pq.Len() > k {
heap.Pop(&pq)
}
}
res := make([]int, 0)
for pq.Len() > 0 {
res = append(res, heap.Pop(&pq).(*Item).value)
}
return res
}