Crushing Tech Interviews with the Top K Elements Pattern

To follow up from my previous post, let's dig more into some patterns that revolve around heaps!

When you're given a question that revolves around keeping track of k elements, a good bet is to think about solving the problem with heaps.

Lets say you're given a problem like the following:

Given an unsorted array of numbers, find the k largest numbers in it.

Example:

Input: [3, 1, 5, 13, 2, 12], K = 3
Output: [5, 13, 12]

First that comes to mind? Let's sort it and return the largest k elements.

Easy!

Unfortunately this will yield a time complexity of O(nlogn) and guess what? We can do better.

If we iterate through the list one at a time, we can maintain a list of k elements by adding the current element to a min heap. If the length of the min heap becomes greater than k then all we need to do is pop off the smallest element which is conveniently stored at the top of the heap!

Popping off the smallest element of the heap takes O(logk) complexity due to the reshuffling that needs to happen to re-order the remaining elements in the heap.

Given the example above, we can see the algorithm at work here:

1. We add to the heap until we've got 3 elements.
2. We add 13 to the heap. len(heap) > 3 so we pop off the top element and reshuffle.
3. 2 is smaller than the top element of the heap so we ignore it.
4. We add 12 to the heap. len(heap) > 3 so we pop off the top element and reshuffle.

def find_k_largest(arr: List[int], k: int) -> List[int]:
    min_heap = []

    for num in arr[0:k]:
        heappush(min_heap, num)

    for num in arr[k:]:
        if num > min_heap[0]:
            heappop(min_heap)
            heappush(min_heap, num)

    return list(min_heap)

Here's some extra questions to get you familiar with the pattern:

https://leetcode.com/problems/third-maximum-number/ (Easy)

https://leetcode.com/problems/top-k-frequent-elements/ (Medium)

https://leetcode.com/problems/kth-largest-element-in-an-array/ (Medium)