347. Top K Frequent Elements

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

Constraints:

• 1 <= nums.length <= 105
• -104 <= nums[i] <= 104
• k is in the range [1, the number of unique elements in the array].
• It is guaranteed that the answer is unique.

There are two main ways to solve this problem: the intuitive approach using a heap, and a more efficient method using buckets.

Heap Approach

A heap stores information about how frequent an element appears, and can efficiently return the most frequent while editing the heap, so it seems to be a natural method to approach this problem. Simply build a max heap using a frequency map, and then pop the $k$ most frequent off the heap and return.

def topKFrequent(self, nums: List[int], k: int) -> List[int]:
	import heapq
	
	count = {}
	for num in nums:
		if num not in count:
			count[num] = 1
		else:
			count[num] += 1
	
	heap = [(count[key], key) for key in count]
	heapq.heapify_max(heap)
	
	ans = []
	for i in range(0, k):
		ans.append(heapq.heappop_max(heap)[1])
	
	return ans

Time Complexity

Heapify takes $\mathcal{O}(n)$ time, and so does building count. Each iteration of the loop takes $\mathcal{O}(\log n)$ time due to the heap pop operation. There are $k$ iterations, but notice that $k$ is in $\mathcal{O}(n)$. Thus the loop takes $\mathcal{O}(n \log n )$ run time. Since that dominates, the total algorithm has $\mathcal{O}(n \log n)$ time complexity.

Space Complexity

The heap takes $\mathcal{O}(n)$ space, since in the worst case it stores $n$ different elements. All other variables take constant space, so the total space complexity is in $\mathcal{O}(n)$.

Bucket Sort

Consider two important points when gathering the top $k$ elements.

• The data is fixed—there is no need to constantly update our understanding of what the most frequent elements are. Once elements have been sorted by frequency, we don’t need to do any further editing of the heaps.
• Order doesn’t matter, so we don’t need to gather the top $k$ element in a sorted order—and there is no need to perform any $\mathcal{O}(\log n)$ sorting after gathering.

A heap therefore is not necessarily needed, since we don’t need to maintain an order nor deal with shifting frequencies. Instead, we can simply sort elements by frequency, and then return the $k$ most frequent.

def topKFrequent(nums: List[int], k: int) -> List[int]:
	Build frequency map of elements
	
	sorted = []
	Place elements into sorted according to frequency
	
	res = []
	for i in range(len(sorted), len(sorted) - k, -1):
		res.append(sorted[i])
		
	return res

The question is then how we can sort elements by frequency efficiently. Naively sorting takes $\mathcal{O}(\log n)$ and thus does not improve on the heap algorithm:

sorted_elements = sorted(count.keys, key=lambda num: count[num])
return sorted_elements[-k:]

Instead, notice that a frequency of an element can only range between 1 and $n$, where all the elements of nums are the same. And since we can create a frequency map of the elements efficiently, we can group elements by frequencies into “buckets”. Since order doesn’t matter, we can iterate through these groups however we wish to return the most frequent ones!

def topKFrequent(nums: List[int], k: int) -> List[int]:
	# build frequency map
	count = {}
	for num in nums:
		if num not in count:
			count[num] = 1
		else:
			count[num] += 1
			
	# build "buckets" of frequencies
	sorted_freqs = [set() for _ in range(len(nums))]
	for num in count:
		# for every unique element, put it in 
		# its respective frequency bucket since we are
		# guaranteed that its frequency is between 1 and n
		sorted_freqs[count[num]].add(num)
	
	res = []
	# going from most frequent to least, add elements until
	# there are k elements in the result
	for i in range(len(sorted_freqs) - 1, 0, -1):
		for num in sorted_freqs[i]:
			res.append(num)
			if len(res) == k:
				return res
				
	return res

Time Complexity

In the worst case, there are $n$ unique elements all with frequency $1$. Then building the frequency map takes $\mathcal{O}(n)$ time, building the buckets $\mathcal{O}(n)$ time, and adding elements $\mathcal{\theta}(k)$ time. Since $k \in \mathcal{O}(n)$ as discussed above, the full algorithm is in $\mathcal{O}(n)$ time.

Space Complexity

Again, in the worst case, there are $n$ unique elements all with frequency $1$. The frequency map then has $n$ elements and the array of buckets has $n$ buckets with 1 element each. The total space complexity is therefore in $\mathcal{O}(n)$.