top-k-freq-elements.py

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        # lets use the bucket sort
        # dict
        #   count: [0, 1, 2, 3]
        #   freq: [[], [3], [2], [1]]
        # basically values how many time have appeared 
        # (index is count)

        n = len(nums)

        count = {}
        freq = [[] for i in range(n + 1)]

        for n in nums:
            count[n] = 1 + count.get(n, 0)  
    
        for n, c in count.items():
            freq[c].append(n)

        res = []
        for i in range(len(freq) - 1, 0, -1):
            for n in freq[i]:
                res.append(n)
                if len(res) == k:
                    return res