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

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