Top K Frequent Elements

Min-Heap

Problem

Given a non-empty array of integers, return the k most frequent elements.

For example:

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

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

Thought Process

• This problem is very similar to how we solved K Closest Points but in this problem we are storing (frequency, number) pairs in our heap.

• We want the top K elements, this is why a min-heap is good.

Solution

from heapq import *
class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        freq = {}
        heap = []
        
        for i in nums:
            if i not in freq:
                freq[i] = 0
            freq[i]+=1
        
        for num, count in freq.items():
            heappush(heap, (count, num))
            if len(heap) > k:
                heappop(heap)
        
        res = []
        for i in range(len(heap)):
            res.append(heap[i][1])
        return res
                

Key Points

• Using min-heap to store elements according to their frequency count so we get the top K

• Storing the frequency count and number together in pairs in the heap

Time Complexity

• Time: $O(N+N*log(K))$

• Space: $O(N)$ because even though we are storing only ‘K’ numbers in the heap. For the frequency map, however, we need to store all the ‘N’ numbers.