LeetCode 347 - Top K Frequent Elements - 2023.10.04 solution

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        """ If you had a dict mapping from each number to its count, along
        with a variable containing the currently-most-common number, you could
        step through the array, updating the dict and the variable as you go.
        That would be o(n) time but also O(n) memory in the worst case (where
        every number is different). That would be only to get the most-
        frequent number.  If we wanted the two-most-frequent, we could use
        two variables, or a list.

        The naive approach would be to just keep a dict of each num to its
        count, then turn it into tuples and sort by the count, and step
        through the sorted version to return the k-most-frequent. That would
        I think be O(nlogn) in the worst case (where each number only appears
        once).

        My thought now is: if we absolutely want O(n) time, how can we throw
        more memory at the problem?
        Another question I ask myself: what could we do with a second pass
        through the input once we'd done our first pass and created that
        dict, or what could we do by doing an iteration through the dict?

        Once we have the dict, we could create another list of the length of
        the number of elements, step through the dict and add each number
        to the index corresponding to its count.  Then we could do an
        iteration backwards through that list to find the top k elements.
        Finished planning / thinking at 14:07
        """
        num_to_count = defaultdict(int)
        for n in nums:
            num_to_count[n] += 1

        count_to_nums = [[] for i in range(len(nums))]
        for n in num_to_count.keys():
            count = num_to_count[n]
            count_to_nums[count-1].append(n)

        top_k_elements = []
        j = 0
        for i in range(len(nums)-1, -1, -1):
            if j == k:
                return top_k_elements
            if len(count_to_nums[i]) > 0:
                top_k_elements.extend(count_to_nums[i])
                j += len(count_to_nums[i])
        return top_k_elements