less than 1 minute read

347. Top K Frequent Elements

Difficulty: Medium

Related Topics: Array, Hash Table, Sorting, Bucket Sort

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]

Code 1:

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        res = []
        count = {}

        for i in nums:
            if i in count:
                count[i] += 1
            else:
                count[i] = 1
        ans = sorted(count.items(), key=lambda x:x[1])
        for i in range(len(ans)-1, len(ans)-1-k, -1):
            res.append(ans[i][0])
        return res

Code 2: Optimized Bucket Sort Solution

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        count  = {}
        table = [[] for _ in range(len(nums)+1)]
        for i in nums:
            count[i] = 1 + count.get(i, 0)
        
        for n, c in count.items():
            table[c].append(n)
        
        res = []
        for n in range(len(table)-1, 0, -1):
            for i in table[n]:
                res.append(i)
                if len(res) == k: 
                    return res