Top K Frequent Elements
Problem Statement
Given an integer array nums
and an integer k
, return the k
most frequent elements within the array.
The test cases are generated such that the answer is unique.
You may return the output in any order.
Example 1:
Input: nums = [1,2,2,3,3,3], k = 2Output: [2,3]
Example 2:
Input: nums = [7,7], k = 1Output: [7]
Constraints:
1 <= nums.length <= 10^4
-1000 <= nums[i] <= 1000
1 <= k <= number of distinct elements in nums
You should aim for a solution with O(n)
time and O(n)
space, where n
is the size of the input array.
Recommended Time and Space Complexity
You should aim for a solution with O(n)
time and O(n)
space, where n
is the size of the input array.
Hint 1
A naive solution would be to count the frequency of each number and then sort the array based on each element’s frequency. After that, we would select the top k
frequent elements. This would be an O(n log n)
solution. Though this solution is acceptable, can you think of a better way?
Hint 2
Can you think of an algorithm which involves grouping numbers based on their frequency?
Hint 3
Use the bucket sort algorithm to create n
buckets, grouping numbers based on their frequencies from 1
to n
. Then, pick the top k
numbers from the buckets, starting from n
down to 1
.
Solution
Sorting
def top_k_frequent(nums: List[int], k: int) -> List[int]: count = {} for num in nums: count[num] = 1 + count.get(num, 0)
arr = [] for num, cnt in count.items(): arr.append([cnt, num]) arr.sort()
res = [] while len(res) < k: res.append(arr.pop()[1]) return res
Time complexity: O(n log n)
Space complexity: O(n)
Heap
import heapq
def top_k_frequent(nums: List[int], k: int) -> List[int]: count = {} for num in nums: count[num] = 1 + count.get(num, 0)
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
Time complexity: O(n log k)
Space complexity: O(n + k)
Bucket Sort
def top_k_frequent(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
Time complexity: O(n)
Space complexity: O(n)