347. Top K Frequent Elements (Python)

Related Topic

Heap.

Description

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

Sample I/O

Example 1

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

Example 2

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

Note

• You may assume k is always valid, 1 ≤ k ≤ number of unique elements.
• Your algorithm’s time complexity must be better than O(n log n), where n is the array’s size.
• It’s guaranteed that the answer is unique, in other words the set of the top k frequent elements is unique.
• You can return the answer in any order.

Methodology

First we need count the appearence for each element store it in dictionary. Then we iterate the dictionary, and push the value-key pair (or count-value pair) into heap, if the the heap length less than k. When heap length is greater or equal than k, we push the current value-key pair into heap and pop one from the heap in order to keep the length of the heap equal to key. After iteration, we will get the k most frequent elements pair, we only need to return the value not the count.

def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        res=[]
        dict = collections.Counter(nums)
        for val, count in dict.items():
            if len(res)<k:
                heapq.heappush(res,(count,val))
            else:
                heapq.heappush(res,(count,val))
                heapq.heappop(res)
        return [val for count, val in res]

BigO

Time complexity : O(n+m+k) which n is the size of input, m is the size of the dictionary and k is given.