Top K Frequent Elements

Problem Statement: 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]

Example 2:

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

Constraints:

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

Source: https://leetcode.com/problems/top-k-frequent-elements/

Leetcode Difficulty: Medium

Asked in: Amazon

Code:

    def topKfrequent(self, nums, k):
        """
        nums: List[int]
        k: int
        return: List[int]
        """
        counter={}
        for i in nums:
            if i in counter:
                counter[i]+=1
            else:
                counter[i]=1
        
        heap=[]
        for idx, (key,value) in enumerate(counter.items()):
            if idx<k:
                heapq.heappush(heap,(value,key))
            else:
                top=heapq.heappop(heap)
                if top[0]<value:
                    heapq.heappush(heap,(value,key))
                else:
                    heapq.heappush(heap,top)
           
        return [heapq.heappop(heap)[1] for _ in range(k)]

Thought Process / Explanation:

The keyword "k frequent" hints me to give a shot to heaps. Heap of size K is enough as we will need to return atmost k elements. Since we need to return k frequent, we can create a min-heap with (value,key) pair and later compare it with the value on the pop operation. Finally, return the key corresponding to each value in min-heap.

Thank You!