Introduction
Heaps are a fundamental data structure to know.
Note that SortedLists are almost strictly more powerful, with little downside. (ie. when you can use a heap, you can use a sorted list. But not necessarily the other way around) However, heaps are more ubiquitous as they are simpler and more commonly taught in school. Familiarity between the two is crucial.
Heaps are complete trees that allow for:
- O(log n) insertions
- O(log n) pop from the top
- O(1) peek top
- O(n) initialization
Compare with sorted lists: O(log n) insertion, O(log n) pop from anywhere, O(log n) search anywhere.
They allow for fast access to min/max elements in an array. There are two types of heaps: min and max. Min heap stores the minimum element at the top, max similarly stores the max at the top.
A trick to get max heaps from min heaps is to negate all values in the heap. Think about why this works: [3,4,5] → [-5,-4,-3] in sorted order. Negating values flips the sorted order, which transforms min → max. Now -5 is at the top of the heap, and if we negate it we get back the original positive 5.
Heap API
Basic operations available on heaps in Python.
Basic Operations
Python's heapq module provides min-heap operations.
from heapq import heapify, heappush, heappop
# heapify(A): makes an array A into a heap, O(n) time
nums = [3, 1, 4, 1, 5]
heapify(nums) # nums is now a valid min-heap
# heappop(heap): pops the top element in O(log n) time, rebalances
min_val = heappop(nums)
# heappush(heap, x): add x to heap in O(log n), rebalances
heappush(nums, 2)
# heap[0]: looks at the top of the heap in O(1)
top = nums[0]Tuple Sorting in Heaps
Sorting in tuples goes first by the smaller first value earlier, then for tiebreakers we move to the next value. If the first values are the same, the one with the smaller second value will be earlier.
# Example: (-value, row) - sorts by largest value first
heap = [(-grid[i][j], i) for i in range(m) for j in range(n)]
heapify(heap)
# When popping, negate again to get original value
val, row = heappop(heap)
original_val = -valBasic Heap Problems
Problems demonstrating fundamental heap usage.
01Maximum Sum With at Most K ElementsMedium
You are given a 2D integer matrix grid of size n x m, an integer array limits of length n, and an integer k. The task is to find the maximum sum of at most k elements from the matrix grid such that:
• The number of elements taken from the ith row of grid does not exceed limits[i].
Return the maximum sum.
def maxSum(self, grid: List[List[int]], limits: List[int], k: int) -> int:
m, n = len(grid), len(grid[0])
heap = [(-grid[i][j], i) for i in range(m) for j in range(n)]
heapify(heap)
res = 0
while k:
val, row = heappop(heap)
if limits[row] == 0:
continue
limits[row] -= 1
res += -val
k -= 1
return resO(n log n)O(n)We use a max heap (negate values), storing pairs of (-value, row_index). Greedily take the max value every time unless we've hit the limit for that row.
Top K / Bottom K Problems
This is a class of very similar problems and solutions. There are 3 ways to solve top K problems:
- Sort -
O(n log n)time - Heap of size k -
O(n log k)time (or sorted list of size k) - Quick select -
O(n)average time,O(n²)worst case.O(n)guaranteed using median of medians.
Important: If you want top K largest, you need a min heap because you want to remove the smallest and keep the top K largest. Similarly for bottom K smallest, you need a max heap.
01Find the Kth Largest Integer in the ArrayMedium
You are given an array of strings nums and an integer k. Each string in nums represents an integer without leading zeros.
Return the string that represents the kth largest integer in nums.
Note: Duplicate numbers should be counted distinctly.
# Solution 1: Sort
def kthLargestNumber(self, nums: List[str], k: int) -> str:
nums.sort(key=lambda x: int(x))
return nums[-k]
# Solution 2: Min heap of size k
def kthLargestNumber(self, nums: List[str], k: int) -> str:
heap = []
for num in nums:
heappush(heap, int(num))
if len(heap) == k + 1:
heappop(heap)
return str(heap[0])O(n log k)O(k)For kth largest, use a min heap of size k. When heap exceeds size k, pop the smallest. The top of the heap is the kth largest.
02Kth Largest Element in an ArrayMedium
Given an integer array nums and an integer k, return the kth largest element in the array.
Note that it is the kth largest element in the sorted order, not the kth distinct element.
Can you solve it without sorting?
# Solution 1: Sort
def findKthLargest(self, nums: List[int], k: int) -> int:
nums.sort()
return nums[-k]
# Solution 2: Min heap of size k
def findKthLargest(self, nums: List[int], k: int) -> int:
heap = []
for num in nums:
heappush(heap, num)
if len(heap) == k + 1:
heappop(heap)
return heap[0]
# Solution 3: Quick Select (O(n) average)
def findKthLargest(self, nums: List[int], k: int) -> int:
if not nums:
return
pivot = random.choice(nums)
left = [x for x in nums if x > pivot]
mid = [x for x in nums if x == pivot]
right = [x for x in nums if x < pivot]
L, M = len(left), len(mid)
if k <= L:
return self.findKthLargest(left, k)
elif k > L + M:
return self.findKthLargest(right, k - L - M)
else:
return mid[0]O(n) average for quickselectO(n)Quick select partitions the array around a pivot and recurses on only one side based on k's position relative to the partition.
03K Closest Points to OriginMedium
Given an array of points where points[i] = [xi, yi] represents a point on the X-Y plane and an integer k, return the k closest points to the origin (0, 0).
# Solution 1: Sort
def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
# Trick: compare squared euclidean distance (x^2 is monotonic)
# This saves compute because square roots are expensive
points.sort(key=lambda x: x[0]*x[0] + x[1]*x[1])
return points[:k]
# Solution 2: Max heap of size k (for bottom K, use max heap)
def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
heap = []
for x, y in points:
heappush(heap, (-(x*x + y*y), x, y))
if len(heap) == k + 1:
heappop(heap)
return [[x, y] for _, x, y in heap]O(n log k)O(k)We need bottom K (smallest distances), so we use a max heap (negate values). The heap keeps the k closest points.
04Top K Frequent WordsMedium
Given an array of strings words and an integer k, return the k most frequent strings.
Return the answer sorted by the frequency from highest to lowest. Sort the words with the same frequency by their lexicographical order.
# Solution 1: Sort
def topKFrequent(self, words: List[str], k: int) -> List[str]:
freq = Counter(words)
words = sorted((-freq[word], word) for word in set(words))
return [word for _, word in words[:k]]
# Solution 2: Heap (push all, pop k times)
def topKFrequent(self, words: List[str], k: int) -> List[str]:
freq = Counter(words)
heap = []
for word in set(words):
heappush(heap, (-freq[word], word))
res = []
for _ in range(k):
res.append(heappop(heap)[1])
return resO(n log n)O(n)Note: Doing this in O(n log k) with a heap of size k is tricky because of the lexicographical tiebreaker requirement.
2 Heap Solutions (Median)
The idea is we want to quickly find the median. We do this by storing 2 heaps:
- One max heap on the left with equal size or 1 larger than the right heap
- One min heap on the right
Now:
- If there is an odd number of elements, the median is the top of the left heap
- If there is an even number of elements, the median is the average of the two tops
Two invariants:
1. Size invariant: len(left) >= len(right) and len(left) <= len(right) + 1
2. Value invariant: top of left <= top of right
Add Operations Walkthrough
stream = [2, 1, 7, 9]
- Add 2: [2][] → [][2] → [2][]
- Add 1: [1,2][] → [1][2]
- Add 7: [1,7][2] → [1][2,7] → [1,2][7]
- Add 9: [1,2,9][7] → [1,2][7,9]
Add Algorithm
- Always add to the left heap first
- Move the top (max) of the left heap to the right heap
- If right has more elements than left, move top of right to left
- If left has more than right+1 elements, move top of left to right (for removals)
def addNum(self, num: int) -> None:
# Push to max-heap (invert to simulate max-heap)
heapq.heappush(self.left, -num)
# Balance step 1: move largest from left to right
heapq.heappush(self.right, -heapq.heappop(self.left))
# Balance step 2: if right has more, move smallest back to left
if len(self.left) < len(self.right):
heapq.heappush(self.left, -heapq.heappop(self.right))Remove Algorithm (requires SortedList)
Heaps don't support O(log n) deletion, only O(n). For removals, use SortedList instead.
def _rebalance(self):
# move the smallest from right to left
if len(self.left) < len(self.right):
self.left.add(self.right.pop(0))
# move the largest from left to right
elif len(self.left) > len(self.right) + 1:
self.right.add(self.left.pop(-1))
def removeNum(self, num: int) -> None:
# try removing from left; if it fails, remove from right
try:
self.left.remove(num)
except ValueError:
self.right.remove(num)
# after removal, rebalance
self._rebalance()01Find Median from Data StreamHard
The median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value, and the median is the mean of the two middle values.
Implement the MedianFinder class:
• MedianFinder() initializes the MedianFinder object.
• void addNum(int num) adds the integer num from the data stream.
• double findMedian() returns the median of all elements so far.
# Solution 1: Two Heaps
class MedianFinder:
def __init__(self):
self.left = [] # Max-heap (invert values)
self.right = [] # Min-heap
def addNum(self, num: int) -> None:
heapq.heappush(self.left, -num)
heapq.heappush(self.right, -heapq.heappop(self.left))
if len(self.left) < len(self.right):
heapq.heappush(self.left, -heapq.heappop(self.right))
def findMedian(self) -> float:
if len(self.left) > len(self.right):
return -self.left[0]
return (-self.left[0] + self.right[0]) / 2
# Solution 2: SortedList (trivializes the problem)
class MedianFinder:
def __init__(self):
self.arr = SortedList()
def addNum(self, num: int) -> None:
self.arr.add(num)
def findMedian(self) -> float:
n = len(self.arr)
if n % 2 == 1:
return self.arr[n // 2]
return (self.arr[n // 2] + self.arr[n // 2 - 1]) / 2O(log n) per operationO(n)Using a single SortedList completely trivializes this problem, but the two-heap solution is more commonly expected in interviews.
02Sliding Window MedianHard
You are given an integer array nums and an integer k. There is a sliding window of size k which is moving from left to right. Return the median array for each window.
from sortedcontainers import SortedList
class MedianFinder:
def __init__(self):
self.left = SortedList() # smaller half (or one extra)
self.right = SortedList() # larger half
def _rebalance(self):
if len(self.left) < len(self.right):
self.left.add(self.right.pop(0))
elif len(self.left) > len(self.right) + 1:
self.right.add(self.left.pop(-1))
def addNum(self, num: int) -> None:
self.left.add(num)
self.right.add(self.left.pop(-1))
self._rebalance()
def removeNum(self, num: int) -> None:
try:
self.left.remove(num)
except ValueError:
self.right.remove(num)
self._rebalance()
def findMedian(self) -> float:
if len(self.left) > len(self.right):
return float(self.left[-1])
return (self.left[-1] + self.right[0]) / 2
def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
mf = MedianFinder()
n = len(nums)
res = []
for i in range(n):
mf.addNum(nums[i])
if i >= k - 1:
res.append(mf.findMedian())
mf.removeNum(nums[i - (k - 1)])
return resO(n log k)O(k)Uses 2 SortedLists instead of 2 heaps because heaps do not support O(log n) deletion. For a MedianFinder that needs removals, we need SortedLists.
03Finding MK AverageHard
Design data structure that calculates MKAverage: average of last m elements excluding k smallest and k largest.
from sortedcontainers import SortedList
class MKAverage:
def __init__(self, m: int, k: int):
self.m, self.k = m, k
self.left = SortedList() # k smallest
self.middle = SortedList() # m - 2k middle elements
self.right = SortedList() # k largest
self.middle_sum = 0
self.stream = []
def addElement(self, num: int) -> None:
self.stream.append(num)
# Remove element if window exceeds m
if len(self.stream) > self.m:
rm = self.stream[-self.m - 1]
if rm in self.left:
self.left.remove(rm)
elif rm in self.middle:
self.middle.remove(rm)
self.middle_sum -= rm
else:
self.right.remove(rm)
# Add to left, bubble up through middle to right
self.left.add(num)
if len(self.left) > self.k:
val = self.left.pop(-1)
self.middle.add(val)
self.middle_sum += val
if len(self.middle) > self.m - 2 * self.k:
val = self.middle.pop(-1)
self.middle_sum -= val
self.right.add(val)
# Rebalance: ensure left has k, middle has m-2k
while len(self.left) < self.k and self.middle:
val = self.middle.pop(0)
self.middle_sum -= val
self.left.add(val)
while len(self.middle) < self.m - 2 * self.k and self.right:
val = self.right.pop(0)
self.middle.add(val)
self.middle_sum += val
def calculateMKAverage(self) -> int:
if len(self.stream) < self.m:
return -1
return self.middle_sum // len(self.middle)O(log m) per operationO(m)Extension of median finding with 3 SortedLists. Maintain k smallest in left, k largest in right, middle contains the rest. Track middle_sum for O(1) average calculation.
04Find X-Sum of All K-Long Subarrays IIHard
For each k-length window, find sum of x most frequent elements (by frequency, then by value).
from sortedcontainers import SortedList
from collections import Counter
def findXSum(self, nums: List[int], k: int, x: int) -> List[int]:
res = []
left = SortedList() # (freq, val) not in top x
right = SortedList() # (freq, val) in top x
sum_right = 0
freq = Counter()
def add(val):
nonlocal sum_right
f = freq[val]
if f > 0:
if (f, val) in right:
right.remove((f, val))
sum_right -= f * val
else:
left.discard((f, val))
freq[val] += 1
f += 1
left.add((f, val))
# Move from left to right if needed
while len(right) < x and left:
item = left.pop(-1)
right.add(item)
sum_right += item[0] * item[1]
# Swap if left has larger than right's smallest
while left and right and left[-1] > right[0]:
l_item = left.pop(-1)
r_item = right.pop(0)
left.add(r_item)
right.add(l_item)
sum_right += l_item[0] * l_item[1] - r_item[0] * r_item[1]
def remove(val):
nonlocal sum_right
f = freq[val]
if (f, val) in right:
right.remove((f, val))
sum_right -= f * val
else:
left.remove((f, val))
freq[val] -= 1
if freq[val] > 0:
left.add((freq[val], val))
# Rebalance
while len(right) < x and left:
item = left.pop(-1)
right.add(item)
sum_right += item[0] * item[1]
for i, num in enumerate(nums):
add(num)
if i >= k - 1:
res.append(sum_right)
remove(nums[i - k + 1])
return resO(n log k)O(k)Sliding window with two SortedLists tracking (frequency, value) pairs. Right list holds top x by frequency. Maintain sum of right for O(1) query.
Appendix: QuickSelect
QuickSelect finds the kth smallest element in an array.
Partition rearranges the array such that everything to the left of pivotIdx is <= nums[pivotIdx], and everything to the right is > nums[pivotIdx].
We recurse on the side of the pivotIdx depending on the comparison of pivotIdx and k-1. This partial sorting behavior avoids the O(n log n) cost, costing O(n) average but O(n²) worst case if partition performance is poor (pivotIdx never partitions evenly).
Use median of medians for guaranteed O(n) worst case.
Partition Function
Returns the pivotIdx where everything left is <= pivot and everything right is > pivot.
def partition(arr: List[int], left: int, right: int) -> int:
pivot = arr[right]
i = left
for j in range(left, right):
if arr[j] <= pivot:
arr[i], arr[j] = arr[j], arr[i]
i += 1
arr[i], arr[right] = arr[right], arr[i]
return iQuickSelect Implementation
Both recursive and iterative versions (k is 1-based).
# Recursive
def kthSmallestRecursive(arr: List[int], left: int, right: int, k: int) -> int:
pivotIdx = partition(arr, left, right)
if pivotIdx == k - 1:
return arr[pivotIdx]
elif pivotIdx > k - 1:
return kthSmallestRecursive(arr, left, pivotIdx - 1, k)
else:
return kthSmallestRecursive(arr, pivotIdx + 1, right, k)
# Iterative
def kthSmallest(arr: List[int], k: int) -> int:
left, right = 0, len(arr) - 1
while left <= right:
pivotIdx = partition(arr, left, right)
if pivotIdx == k - 1:
return arr[pivotIdx]
elif pivotIdx > k - 1:
right = pivotIdx - 1
else:
left = pivotIdx + 1
return -1