347. Top K Frequent Elements

LeetCode 347. Top K Frequent Elements

Description

Given an integer array nums and an integer k, return the k most frequent elements. You may return the answer in any order.

Example 1:

Input: nums = [1,1,1,2,2,3], k = 2
Output: [1,2]

Constraints:

• 1 <= nums.legth <= 105

• k is in the range [1, the number of unique elements in the array].

• It is guaranteed that the answer is unique.

Follow up: Your algorithm's time complexity must be better than O(n log n), where n is the array's size.

Tags

Hash Table, Heap

Solution

Define a structure pair to store element and its frequency in nums. Then convert nums into a map whose key is the element and value is the frequency, and convert the each item into a pair. Push every pair into the pairHeap.

Refer to Reference for the following procedure.

Complexity

• Time complexity: $O(n\log(k))$

• Space complexity: $O(n)$, n for the hash table and k for the heap.

Code

type pair struct {
	ele, freq int
}

type pairHeap []pair

func (h pairHeap) Len() int {
	return len(h)
}

func (h pairHeap) Less(i, j int) bool {
	return h[i].freq < h[j].freq
}

func (h pairHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

func (h *pairHeap) Push(x interface{}) {
	*h = append(*h, x.(pair))
}

func (h *pairHeap) Pop() interface{} {
	x := (*h)[len(*h)-1]
	*h = (*h)[:len(*h)-1]
	return x
}

func topKFrequent(nums []int, k int) []int {
	ans, ele2freq, h := make([]int, 0, k), map[int]int{}, make(pairHeap, 0, k)
	for _, num := range nums {
		ele2freq[num]++
	}
	for ele, freq := range ele2freq {
		if len(h) < k {
			heap.Push(&h, pair{ele, freq})
		} else if freq > h[0].freq {
			heap.Pop(&h)
			heap.Push(&h, pair{ele, freq})
		}
	}
	for _, p := range h {
		ans = append(ans, p.ele)
	}
	return ans
}

Reference

1. 215. Kth Largest Element in an Array

2. 703. Kth Largest Element in a Stream