Top K Frequent Elements — Day 59(Python)

Photo by Chinh Le Duc on Unsplash

Today’s question is a common interview question among Captial One interviewers. It is a medium tagged question in Leetcode.

347. Top K Frequent Elements

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

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.

As the first step, we would require a dictionary that can store the frequency of each number in the input list.

dict_numbers = collections.defaultdict(int)
for n in nums:
    dict_numbers[n] += 1

Next, we need to sort the numbers according to the frequency of their occurrence. We will be using the heap package to sort and get frequent K elements.

We will be creating a list out of the dictionary i.e. key and its value, doing this step will help us to create a heap and then sort it using the heapify function.

We need to keep in mind that heappop function will give us the smallest element in our heap. Hence when we are creating a list, the frequency has to be converted to its negative form.

Let us look into the code snippet.

import heapq
class topKFrequentFinder:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        output = []
        dict_numbers = collections.defaultdict(int)
        for n in nums:
            dict_numbers[n] += 1
        heap_list = [(-freq, key) for key, freq in dict_numbers.items()]
        heapq.heapify(heap_list)
        while(len(output) < k):
            output.append(heapq.heappop(heap_list)[1])
        return output

Complexity analysis.

Time Complexity

We are traversing through the list once, to calculate the frequency. That takes O(N). Post that, we are using heap to sort and find frequent words. That takes O(NlogK), where N is the number of items in the list. Hence the time complexity is O(NlogK).

Space Complexity.

We are creating a dictionary to store the elements and their frequency. Hence the space complexity is O(N).