Problem Description

Given an integer array nums and an integer k, return the k most frequent elements. You may return the answer in any order.

Example 1:

  • Input: nums = [1,1,1,2,2,3], k = 2
  • Output: [1,2]

Example 2:

  • Input: nums = [1], k = 1
  • Output: [1]

My Idea

My initial idea was to store the elements and their respective frequencies in a dictionary, whose keys i sort by frequency and return the first k of them. This results in a time complexity of O(nlog(n)) due to the sorting. My next idea was to replace the sorting with a buckets array top, where top[i] contains a list of all elements with frequency i. Then i remove the empty lists from top and iterate over it in reverse to return an array of the last k elements.

My solution

#Time complexity of O(n)
def topKFrequent(nums: list[int], k: int) -> list[int]:
    freqs = dict()
    for n in nums:
        if n in freqs:
            freqs[n]+=1
        else:
            freqs[n]=1
    top=[[]for _ in range(len(nums)+1)]
    for n,f in freqs.items():
        top[f].append(n)
    cleantop = [t for t in top if t!=[]]
    res=[]
    for i in range(len(cleantop)-1,-1,-1):
        for j in cleantop[i]:
            res.append(j)
            if len(res)==k:
                return res

#Initial solution with O(nlog(n))
def topKFrequentInitial(nums: list[int], k: int) -> list[int]:
    freqs = dict()
    for n in nums:
        if n in freqs:
            freqs[n]+=1
        else:
            freqs[n]=1
    freqs = sorted(freqs, key=freqs.get, reverse=True)
    return freqs[:k]