Lesson 37 — Heap I: the best remaining

~20 minutes · Days 45-46 · heap top-k · Python

Tonight's problem is Kth Largest Element in an Array. Interviewers like it because the obvious answer is tempting: sort everything, then index the answer. The heap idea is sharper. If the question only cares about the best few values, keep only those few values alive.

The invariant

A heap hands you the best remaining element in O(log n) per update; keep a min-heap of size k to hold the k largest.

That's the whole pattern. A heap is a structure built so its top element is always the current best according to its ordering. Python's heapq module gives you a min-heap, so heap[0] is the smallest item in the heap. Binary heaps support updates in logarithmic time (Wikipedia: binary heap).

A Python heap is an array, but the useful mental model is the tree root: the smallest value is always easiest to remove.

Kth Largest: the slow way first

The LeetCode statement asks for the kth largest value, not the full sorted array. Sorting works, but it does more ordering than the question needs:

def findKthLargest(nums, k):
    nums = sorted(nums, reverse=True)   # order every value
    return nums[k - 1]                  # then take one value

That is O(n log n) time. The heap version keeps a smaller promise: after each number, the heap contains the k largest values seen so far. If a new value makes the heap too large, pop the smallest candidate out.

Kth Largest: the heap way

import heapq

def findKthLargest(nums, k):
    heap = []                         # the k largest seen so far
    for x in nums:
        heapq.heappush(heap, x)       # x gets a chance
        if len(heap) > k:
            heapq.heappop(heap)       # evict the smallest candidate
    return heap[0]                    # smallest of top k = kth largest

The line to trust is return heap[0]. Because the heap holds exactly the k largest values, its smallest value is the kth largest overall.

The heap keeps the winners by removing the weakest candidate whenever the heap grows past k.

Transfer: same pattern, new costume

Top K Frequent Elements: return the k values that appear most often.

Before opening the reveal: predict what the heap should hold. The input values themselves are not the ranking key anymore. What is?

from collections import Counter
import heapq

def topKFrequent(nums, k):
    heap = []                              # smallest frequency at top
    for value, count in Counter(nums).items():
        heapq.heappush(heap, (count, value))
        if len(heap) > k:
            heapq.heappop(heap)
    return [value for count, value in heap]

What the trade buys you

The heap wins when k is much smaller than n. More importantly, it gives you the interview habit: ask what belongs in the heap, and whether the top should be the next answer or the next eviction.

Practice — answer before you scroll past

for x in nums:
    heapq.heappush(heap, -x)
    if len(heap) > k:
        heapq.heappop(heap)
return -heap[0]

Your move (tonight)

Solve these on LeetCode, in order, logging each in your Notion tracker with the cycle — predict → attempt → compare → encode:

1. Kth Largest Element in an Array — predict the sorting brute force first, then write the size-k heap from memory.
2. Top K Frequent Elements — answer the heap self-check gate before coding: min-heap or max-heap, and what lives in it?

For the encode step, write the invariant at the top of this page in your own words. The important sentence is not the code; it is why the heap top is exactly the next value to remove or return.