LC 347 - Top K Frequent Elements

Question

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
• -104 <= nums[i] <= 104
• 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.

Links

Question here and solution here

Solution

concept

We will first create the counter dictionary that keeps track the count of each number’s appearances.

From here there are two ways:
1) max heap
2) bucket sort

The max heap method basically add all the (-count, num) pair into the heap and then just pop the top k values.
The bucket sort is to keep track an array which is a list of list [[the numbers that have freq 0], [the numbers that have freq 1]]. That the index of this list is the number of count in the first counter dictionary.
We then iterate this list in reverse to get the top k count’s number (through the index of the list).

code

class Solution:
	"""
	max-heap solution
	"""
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        counter = defaultdict(int)
        ans = []

        for num in nums:
            counter[num] += 1

        heap = []
        for num, freq in counter.items():
            heapq.heappush(heap, (-freq, num))

        while k > 0:
            freq, num = heapq.heappop(heap)
            ans.append(num)
            k -= 1

        return ans

class NeetSolution:
	"""
	bucket sort solution
	"""
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        count = {}
        freq = [[] for i in range(len(nums) + 1)]

        for num in nums:
            count[num] = 1 + count.get(num, 0)
        for num, cnt in count.items():
            freq[cnt].append(num)
        
        res = []
        for i in range(len(freq) - 1, 0, -1):
            for num in freq[i]:
                res.append(num)
                if len(res) == k:
                    return res

class NeetSolution:
	"""
	min-heap solution
	"""
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        count = {}
        for num in nums:
            count[num] = 1 + count.get(num, 0)
            
		# cleverly only keep track the top k, since smallest will be pop out
        heap = []
        for num in count.keys():
            heapq.heappush(heap, (count[num], num))
            if len(heap) > k:
                heapq.heappop(heap)

        res = []
        for i in range(k):
            res.append(heapq.heappop(heap)[1])
        return res

Complexity

heap:
time: $O(n \log k)$,
space: $O(n + k)$
where $n$ is the length of the array and $k$ is the number of top frequent elements.

bucket sort:
time: $O(n)$, since the we cleverly choose the number of counts as the one we iterate through as index
space: $O(n)$