leetcode/0347_top-k-frequent-elements/python3/bucketsort.py

42 lines
1.5 KiB
Python

class Solution:
def topKFrequent(self, nums: List[int], k: int) -> List[int]:
'''
Approach using bucket sort
'''
c = Counter(nums)
# Idea here is that the list indices will represent
# the actual counts and we store elements matching thise
# counts under the index. If size n list is provided, max
# count possible is n but lists are 0-indexed so we set
# num_bucket to have n+1 size.
num_bucket = [[] for _ in range(len(nums) + 1)]
# Later down, we will iterate in reverse through `num_bucket` and
# instead of iterating through the entire list, we can take note
# of the max count we observe and start from there instead.
max_count_seen = -1
for num, count in c.items():
num_bucket[count].append(num)
max_count_seen = max(max_count_seen, count)
result = []
j = max_count_seen
# Iterate reverse through num_bucket (note the max_count_seen
# optimization) and fill the result array until k elements are
# added
while j >= 0:
for num in num_bucket[j]:
result.append(num)
# Possible that we filled k elements already in this inner
# loop, so get out
if len(result) == k:
return result
j -= 1