Teaching Kids Programming – Top K Frequent Elements (Heap and Counter)

Teaching Kids Programming: Videos on Data Structures and Algorithms

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]

Constraints:
1 <= nums.length <= 10^5
k is in the range [1, the number of unique elements in the array].
It is guaranteed that the answer is unique.

Follow up: Your algorithm’s time complexity must be better than O(n log n), where n is the array’s size.

Top K Frequent Elements via Counter’s most_common

The Counter returns a dictionary (hash map) that stores the key-value pairs (items) where keys are numbers/elements and values are the corresponding frequencies. We can returned the most common items by using the most_common method.

1
2
3
4

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        c = Counter(nums)
        return [key for key, val in c.most_common(k)]

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        c = Counter(nums)
        return [key for key, val in c.most_common(k)]

It does not provide a least common, but we can workaround by most_common (sorted in the descending order of the frequencies):

1

c.most_common()[-1]

c.most_common()[-1]

Top K Frequent Elements via Sorting by the Counter’s items

The most_common is essential sorting the items by the values (frequencies).

1
2
3
4
5

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        c = Counter(nums)
        a = sorted(c.items(), key=lambda p: p[1], reverse=True)
        return [i for i, j in a[:k]]

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        c = Counter(nums)
        a = sorted(c.items(), key=lambda p: p[1], reverse=True)
        return [i for i, j in a[:k]]

Top K Frequent Elements via Heap

The heapq in Python is a Min Heap (the heap.pop returns the smallest element). We can negate the numbers to make it a max heap. The algorithm pushes all the items in the form of -frequency, num tuple, then extracts k times to retrieve the Top K Frequent elements.

1
2
3
4
5
6
7
8
9
10
11

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        c = Counter(nums)
        a = []
        for num, freq in c.items():
            heappush(a, (-freq, num))
        ans = []
        for _ in range(k):
            _, num = heappop(a)
            ans.append(num)
        return ans        

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        c = Counter(nums)
        a = []
        for num, freq in c.items():
            heappush(a, (-freq, num))
        ans = []
        for _ in range(k):
            _, num = heappop(a)
            ans.append(num)
        return ans        

We can use the nlargest:

1
2
3
4
5
6
7

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        c = Counter(nums)
        a = []
        for num, freq in c.items():
            heappush(a, (freq, num))
        return [j for i, j in heapq.nlargest(k, a, key=itemgetter(0))]

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        c = Counter(nums)
        a = []
        for num, freq in c.items():
            heappush(a, (freq, num))
        return [j for i, j in heapq.nlargest(k, a, key=itemgetter(0))]

or nsmallest:

1
2
3
4
5
6
7

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        c = Counter(nums)
        a = []
        for num, freq in c.items():
            heappush(a, (-freq, num))
        return [j for i, j in heapq.nsmallest(k, a, key=itemgetter(0))]

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        c = Counter(nums)
        a = []
        for num, freq in c.items():
            heappush(a, (-freq, num))
        return [j for i, j in heapq.nsmallest(k, a, key=itemgetter(0))]

–EOF (The Ultimate Computing & Technology Blog) —