前 K 个高频元素

黎 浩然/ 17 12 月, 2021/ 算法/ALGORITHMS/ 0 comments

本题主要使用 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
}

Share this Post

Leave a Comment

您的邮箱地址不会被公开。 必填项已用 * 标注

*
*