"""
https://leetcode.com/problems/top-k-frequent-elements/
Given an integer array nums and an integer k, return the k most frequent elements. You may return the answer in any order.
"""
from collections import defaultdict
from typing import List
def top_k_frequent(nums: List[int], k: int) -> List[int]:
# unique = set(nums)
# assert k <= unique
hash = defaultdict(list)
# freq = [[] for i in range(len(nums) + 1)]
freq = defaultdict(list)
for i in nums:
hash[i] = 1 + hash.get(i, 0)
# assert len(hash) = len(set(nums))
for element, count in hash.items():
freq[count].append(element)
# test = []
# assert len(list(map(test.extend, freq.values()))) == len(set(nums)))
result = []
# prev = max(freq.keys())
for i in sorted(freq.keys(), reverse=True):
# for i in range(len(freq) - 1, 0, -1):
# assert prev >= i
# prev = i
for num in freq[i]:
result.append(num)
if len(result) == k:
return result
raise RuntimeError(f"k ({k}) might be larger than unique elements in nums ({nums})")
Tests
from publicmatt.leet.problems.top_k_frequent_elements import top_k_frequent
def test_simple():
nums = [1, 1, 1, 2, 2, 3]
k = 2
ans = [1, 2]
assert top_k_frequent(nums, k) == ans
def test_small_k():
nums = [1]
k = 1
ans = [1]
assert top_k_frequent(nums, k) == ans