LC347 - 前 K 个高频元素

题目

LeetCode347. 题目链接

分析

第一反应是排序然后取K个就OK了，排序需要对整个数组进行排序，但其实只需要维护K个有序的序列就可以了。分析大(小)顶堆的特性，我选择小顶堆。保证堆的大小为K，这样每加入一个元素，就会弹出一个最小的元素，那么留下在堆中的肯定是当前最大的K个。
另外：如果使用Java可以直接使用PriorityQueue，这是直接封装好的。

LeetCode官方代码

func topKFrequent(nums []int, k int) []int {
    occurrences := map[int]int{}
    for _, num := range nums {
        occurrences[num]++
    }
    h := &IHeap{}
    heap.Init(h)
    for key, value := range occurrences {
        heap.Push(h, [2]int{key, value})
        if h.Len() > k {
            heap.Pop(h)
        }
    }
    ret := make([]int, k)
    for i := 0; i < k; i++ {
        ret[k - i - 1] = heap.Pop(h).([2]int)[0]
    }
    return ret
}

type IHeap [][2]int

func (h IHeap) Len() int           { return len(h) }
func (h IHeap) Less(i, j int) bool { return h[i][1] < h[j][1] }
func (h IHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *IHeap) Push(x interface{}) {
    *h = append(*h, x.([2]int))
}

func (h *IHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

我的代码(有问题)

package main

import (
	"container/heap"
)

func topKFrequent(nums []int, k int) []int {
	mp := map[int]int{}
	// map统计元素以及频率
	for i := 0; i < len(nums);i ++ {
		mp[nums[i]] ++
	}

	pq := &PriorityQueue{}
	heap.Init(pq)
	i := 0
	// 遍历map
	for k, v := range mp {
		heap.Push(pq, Item{
			value:    k,
			priority: v,
			index:    i,
		})
		// 如果堆的大小超过K，就要弹出一个元素，这样就能保证堆中是目前K个最大的
		if pq.Len() > k {
			heap.Pop(pq)
		}
	}

	// 结果数组， 因为是小顶堆，所以倒叙装入数组中
	result := make([]int, k)
	for i = 0; i < k;i ++ {
		result[k - i - 1] = heap.Pop(pq).(Item).value
	}

	return result
}


type Item struct {
	value    int // map<k> 元素
	priority int    // map<v> 元素出现的频率
	index int // 下标
}

type PriorityQueue []Item

func (pq PriorityQueue) Len() int { return len(pq) }

func (pq PriorityQueue) Less(i, j int) bool {
	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]
	*pq = old[0 : n-1]
	return item
}

func main() {
	result := topKFrequent([]int{1, 1, 1, 2, 2, 3}, 2)
	// 输出一下结果数组
	for _, item := range result {
		println(item)
	}
}

问题！

测试用例为示例1，问题是程序运行多次会出现result结果不一致，有时候输出2，3，有时候输出1，2。
我的代码参考golang源码container/heap/example_pq_test.go优先级队列示例实现。