LeetCode 347 - Top K Frequent Elements - 2022.12.29 solution

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        # How to solve it: we could add the elements to a dict, then sort by the counts
        # and return the top k ones, but that would be O(n log n) in the worst case.
        # To do better, we can create a list the size of the input list, add the items
        # from the count dict to the position in the list corresponding to the count,
        # and then go through the list starting from the right.  That should be O(n)
        # in the worst case.
        count_dict = defaultdict(int)
        for num in nums:
            count_dict[num] += 1

        counts_list = [[] for i in range(len(nums) + 1)]
        for num, count in count_dict.items():
            counts_list[count].append(num)

        res = []
        for i in range(len(counts_list)-1, -1, -1):
            if len(res) == k:
                break
            else:
                res.extend(counts_list[i])
        return res