Heap / Priority Queue

Python's heapq is a min-heap. For max-heap, negate values. Push/pop in O(log n). Use for: K largest/smallest, streaming median, task scheduling.

// K Largest Rule: Use a MIN-heap of size K

Counter-intuitive: for K largest, maintain a min-heap of size K. The min-heap's root is the Kth largest. If new element > root, pop root and push new.

Medium K Closest Points to Origin

import heapq

def kClosest(points, k):
    # Sort by distance^2 (no sqrt needed for comparison)
    heap = [(x*x + y*y, x, y) for x, y in points]
    heapq.heapify(heap)
    return [[x, y] for _, x, y in heapq.nsmallest(k, heap)]
# Time: O(n log k) | Space: O(n)

Medium Task Scheduler

INSIGHT: Max-heap + cooldown queue. Always schedule the most frequent available task. If no task is available, idle. Use a queue to track when tasks become available again.

from collections import Counter, deque
import heapq

def leastInterval(tasks, n):
    count = Counter(tasks)
    max_heap = [-c for c in count.values()]
    heapq.heapify(max_heap)
    cooldown = deque()   # (available_at_time, -count)
    time = 0
    while max_heap or cooldown:
        time += 1
        if max_heap:
            remaining = 1 + heapq.heappop(max_heap)  # decrement count
            if remaining < 0:
                cooldown.append((time + n, remaining))
        if cooldown and cooldown[0][0] == time:
            heapq.heappush(max_heap, cooldown.popleft()[1])
    return time

Hard Find Median From Data Stream ⭐

INSIGHT: Two heaps: max-heap (lower half) + min-heap (upper half). Balance sizes so lower half has at most 1 more element than upper. Median = top of lower half or average of both tops.

import heapq

class MedianFinder:
    def __init__(self):
        self.small = []   # max-heap (negate) — lower half
        self.large = []   # min-heap — upper half

    def addNum(self, num):
        heapq.heappush(self.small, -num)
        # Ensure small's max <= large's min
        if self.small and self.large and -self.small[0] > self.large[0]:
            heapq.heappush(self.large, -heapq.heappop(self.small))
        # Balance sizes: small can have at most 1 more
        if len(self.small) > len(self.large) + 1:
            heapq.heappush(self.large, -heapq.heappop(self.small))
        if len(self.large) > len(self.small):
            heapq.heappush(self.small, -heapq.heappop(self.large))

    def findMedian(self):
        if len(self.small) > len(self.large):
            return -self.small[0]
        return (-self.small[0] + self.large[0]) / 2
# addNum: O(log n) | findMedian: O(1)

← Trees Dynamic Programming →