Top K Frequent Elements

Patrick Leaton
Published: June 6, 2021, 11:43 a.m.
Problem Description

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.

 

Follow up: Your algorithm's time complexity must be better than O(n log n), where n is the array's size.

The description was taken from https://leetcode.com/problems/top-k-frequent-elements/.

Problem Solution

#O(N(Log K)) Time, O(N) Space
class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        heap, output, seen = [], [], {}
       
        for num in nums:
            if num not in seen:
                seen[num] = -1
            else:
                seen[num] -= 1
       
        for num in seen:
            heapq.heappush(heap, (seen[num], num))
           
        while heap and k > 0:
            count, top_frequent = heapq.heappop(heap)
            output.append(top_frequent)
            k -= 1
 

        return output

Problem Explanation

Right off the bat, the problem title states we are returning the k greatest frequent elements.

The "k" is usually an indication that this can be solved by utilizing a heap because it implies that we will be working off some kind of sorted order.  The "greatest" tells us that this may be a max heap.

From those clues, the implementation is how we may expect.  We will use a hashmap to count the frequency of each number in the input array then we will take those frequencies and push them into a heap.

We will then continuously pop the most frequent number off of the heap and append it to the output.  Once we have done that k times, we are done.

Let's start by initializing our heap, our output, and our seen hashmap.

        heap, output, seen = [], [], {}

 

Next, we will go through the array and count the times we see each element.

        for num in nums:

 

If the element isn't in the hashmap already, we will put it there with an initial count of negative one.

            if num not in seen:
                seen[num] = -1

 

If it is in the hashmap, we will decrement the count by one.

            else:
                seen[num] -= 1

 

To create a make-shift max heap from a min-heap, we are making the heap input, the frequency, negative so that the largest frequency in the heap becomes the smallest when negative.

After have counted frequencies, we will go through each element in the seen hashmap and push the frequency and itself into the heap.  It is important the frequency is first because the heap sorts on the first value inside the tuple.

        for num in seen:
            heapq.heappush(heap, (seen[num], num))

 

Then, while we still have values in the heap and while we haven't found the k most frequent elements, we will pop the top frequent element off of the heap along with its count.

        while heap and k > 0:
            count, top_frequent = heapq.heappop(heap)

 

We will append the current top frequent element to the output and decrement k since we just found one more frequent element.

            output.append(top_frequent)
            k -= 1

 

Once we have run out of frequencies or have found k most frequent numbers, we will return those numbers in the output.

        return output