leetcode-go/leetcode/medium/347.top-k-frequent-elements.go

type pair struct {
	_1 int
	_2 int
}

type pairs []pair

func (ps pairs) Len() int           { return len(ps) }
func (ps pairs) Less(i, j int) bool { return ps[i]._1 < ps[j]._1 }
func (ps pairs) Swap(i, j int)      { ps[i], ps[j] = ps[j], ps[i] }

func (ps *pairs) Push(x interface{}) {
	*ps = append(*ps, x.(pair))
}

func (ps *pairs) Pop() interface{} {
	i := ps.Len() - 1
	x := (*ps)[i]
	*ps = (*ps)[:i]
	return x
}

func topKFrequent(nums []int, k int) []int {
	freq := make(map[int]int)
	for _, i := range nums {
		freq[i]++
	}
	var topK pairs
	i := 0
	for key, val := range freq {
		heap.Push(&topK, pair{val, key})
		if k == i {
			heap.Pop(&topK)
		} else {
			i++
		}
	}
	res := make([]int, k)
	for i := range topK {
		res[i] = topK[i]._2
	}
	return res
}