Heap & Priority Queue: Practice Problems with Solutions

Coding Patterns - Guide to 30 Core Algorithm Patterns

Heap & Priority Queue: Complete Guide with Examples

Welcome to the practice problems for heap and priority queue. If you need a refresher on the code, the code templates have the patterns in all 6 languages. Each problem below includes a hint, a visual walkthrough, and the full solution.

The Heap Helpers for JavaScript and Ruby

Python, Java, C++, and Go have standard heaps, so the solutions below use them directly. JavaScript and Ruby do not, so every JavaScript and Ruby solution on this page uses the Heap helper class from the code templates . In short, new Heap() is a min-heap, and a comparator like (a, b) => a > b (JavaScript) or { |a, b| a > b } (Ruby) turns it into a max-heap. The solutions below only show the problem logic.

Recommended Study Order

Start with Last Stone Weight. It is the smallest possible heap program. You build a max-heap and then keep popping. Kth Largest Element in a Stream adds the size cap, which is the idea behind every top-K problem. Kth Largest Element in an Array applies the same cap to a static input. Top K Frequent Elements adds a second key, frequency, which is where custom comparators start mattering. K Closest Points to Origin compares by distance instead of raw value. Merge k Sorted Lists is the first problem where the heap holds the next candidate from each list instead of a finished value. Find Median from Data Stream ends the set by combining a min-heap and a max-heap into one structure. If you can explain why the two median heaps stay balanced, you understand heaps.

The order above is designed to build intuition progressively. The app schedules your reviews so you do not forget the Easy patterns while grinding Hard ones. Set up your review schedule .

Easy Problems

1. Last Stone Weight

LeetCode 1046 | Difficulty: Easy

Brief: Smash the two heaviest stones together. If they weigh the same, both are destroyed. Otherwise the difference goes back into the pile. Return the weight of the last stone, or 0 if nothing is left.

Why this pattern: The problem repeatedly needs the two largest elements in a collection that keeps changing. A max-heap serves both pops and the occasional re-insert in

O(log N)

each.

Key Insight: Every iteration runs the same three steps. Pop twice, then push the difference if it is non-zero. The heap keeps the heaviest stone at the root, so the two pops never need to scan the pile.

Complexity: Time

O(N log N)

, Space

O(N)

Visual:

    graph TD
    A["Stones: 2, 7, 4, 1, 8, 1"] --> B["Pop 8 and 7, push 1"]
    B --> C["Stones: 2, 4, 1, 1, 1"]
    C --> D["Pop 4 and 2, push 2"]
    D --> E["Stones: 2, 1, 1, 1"]
    E --> F["Pop 2 and 1, push 1"]
    F --> G["Pop 1 and 1, both destroyed"]
    G --> H["Last stone: 1"]

Code:

var lastStoneWeight = function(stones) {
    // max-heap: the root holds the heaviest stone
    const heap = new Heap(stones, (a, b) => a > b);

    while (heap.size() > 1) {
        const heaviest = heap.pop();
        const second = heap.pop();
        if (heaviest !== second) {
            heap.push(heaviest - second);
        }
    }

    return heap.size() === 0 ? 0 : heap.peek();
};

import heapq

def last_stone_weight(stones):
    # heapq is a min-heap, so store negative weights to
    # keep the heaviest stone at the root
    heap = [-w for w in stones]
    heapq.heapify(heap)

    while len(heap) > 1:
        heaviest = -heapq.heappop(heap)
        second = -heapq.heappop(heap)
        if heaviest != second:
            heapq.heappush(heap, -(heaviest - second))

    return -heap[0] if heap else 0

class Solution {
    public int lastStoneWeight(int[] stones) {
        // a max-heap, so the two heaviest stones are always
        // the two pops at the root
        PriorityQueue<Integer> heap = new PriorityQueue<>((a, b) -> b - a);
        for (int w : stones) {
            heap.offer(w);
        }

        while (heap.size() > 1) {
            int heaviest = heap.poll();
            int second = heap.poll();
            if (heaviest != second) {
                heap.offer(heaviest - second);
            }
        }

        return heap.isEmpty() ? 0 : heap.peek();
    }
}

class Solution {
public:
    int lastStoneWeight(vector<int>& stones) {
        // the default priority_queue is a max-heap
        priority_queue<int> heap;
        for (int w : stones) {
            heap.push(w);
        }

        while (heap.size() > 1) {
            int heaviest = heap.top();
            heap.pop();
            int second = heap.top();
            heap.pop();
            if (heaviest != second) {
                heap.push(heaviest - second);
            }
        }

        return heap.empty() ? 0 : heap.top();
    }
};

func lastStoneWeight(stones []int) int {
    h := &IntMaxHeap{}
    heap.Init(h)
    for _, w := range stones {
        heap.Push(h, w)
    }

    for h.Len() > 1 {
        heaviest := heap.Pop(h).(int)
        second := heap.Pop(h).(int)
        if heaviest != second {
            heap.Push(h, heaviest-second)
        }
    }

    if h.Len() == 0 {
        return 0
    }
    return heap.Pop(h).(int)
}

type IntMaxHeap []int

func (h IntMaxHeap) Len() int           { return len(h) }
func (h IntMaxHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h IntMaxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *IntMaxHeap) Push(x interface{}) {
    *h = append(*h, x.(int))
}

func (h *IntMaxHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}

def last_stone_weight(stones)
    # max-heap: the root holds the heaviest stone
    heap = Heap.new(stones) { |a, b| a > b }

    while heap.size > 1
        heaviest = heap.pop
        second = heap.pop
        heap.push(heaviest - second) if heaviest != second
    end

    heap.size == 0 ? 0 : heap.peek
end

The solution is the max-heap template with no twists. Building the heap with a batch operation costs

O(N)

, and each of the at most N smashes costs two pops and one push. The only edge case is the final read. When the heap empties out, the answer is 0, not an error.

2. Kth Largest Element in a Stream

LeetCode 703 | Difficulty: Easy

Brief: Design a class that receives a stream of integers. Every add call inserts one value and returns the kth largest element of everything seen so far.

Why this pattern: The size-k min-heap is the exact shape of this problem. The heap holds the k largest values seen so far, and its root is the answer.

Key Insight: The constructor shrinks the initial heap to size k once. Every add then does push, pop-if-overflow, and read the root. No scanning, no sorting, no memory that grows with the stream.

Complexity: Time

O(log k)

per add, Space

O(k)

Visual:

    graph TD
    A["k = 3, heap keeps size 3"] --> B["Add 3"]
    B --> C["Heap: 4, 5, 8, return 4"]
    C --> D["Add 5"]
    D --> E["Heap: 5, 5, 8, return 5"]
    E --> F["Add 10"]
    F --> G["Heap: 5, 8, 10, return 5"]
    G --> H["Add 9"]
    H --> I["Heap: 8, 9, 10, return 8"]

Code:

var KthLargest = function(k, nums) {
    this.k = k;
    // min-heap capped at k: the root is the kth largest
    this.heap = new Heap(nums);
    while (this.heap.size() > k) {
        this.heap.pop();
    }
};

KthLargest.prototype.add = function(val) {
    this.heap.push(val);
    if (this.heap.size() > this.k) {
        this.heap.pop();
    }
    return this.heap.peek();
};

import heapq

class KthLargest:
    def __init__(self, k, nums):
        self.k = k
        self.heap = nums
        heapq.heapify(self.heap)
        # shrink to size k so the root is the kth largest
        while len(self.heap) > k:
            heapq.heappop(self.heap)

    def add(self, val):
        heapq.heappush(self.heap, val)
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)
        return self.heap[0]

class KthLargest {
    // the default PriorityQueue is a min-heap
    private final PriorityQueue<Integer> heap = new PriorityQueue<>();
    private final int k;

    public KthLargest(int k, int[] nums) {
        this.k = k;
        for (int num : nums) {
            heap.offer(num);
        }
        while (heap.size() > k) {
            heap.poll();
        }
    }

    public int add(int val) {
        heap.offer(val);
        if (heap.size() > k) {
            heap.poll();
        }
        return heap.peek();
    }
}

class KthLargest {
private:
    int k;
    priority_queue<int, vector<int>, greater<int>> heap;

public:
    KthLargest(int k, vector<int>& nums) : k(k) {
        for (int num : nums) {
            heap.push(num);
        }
        while (heap.size() > k) {
            heap.pop();
        }
    }

    int add(int val) {
        heap.push(val);
        if (heap.size() > k) {
            heap.pop();
        }
        return heap.top();
    }
};

type KthLargest struct {
    k    int
    heap *IntMinHeap
}

func Constructor(k int, nums []int) KthLargest {
    h := &IntMinHeap{}
    heap.Init(h)
    for _, num := range nums {
        heap.Push(h, num)
    }
    for h.Len() > k {
        heap.Pop(h)
    }
    return KthLargest{k: k, heap: h}
}

func (this *KthLargest) Add(val int) int {
    heap.Push(this.heap, val)
    if this.heap.Len() > this.k {
        heap.Pop(this.heap)
    }
    return (*this.heap)[0]
}

type IntMinHeap []int

func (h IntMinHeap) Len() int           { return len(h) }
func (h IntMinHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h IntMinHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *IntMinHeap) Push(x interface{}) {
    *h = append(*h, x.(int))
}

func (h *IntMinHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}

class KthLargest
    def initialize(k, nums)
        @k = k
        @heap = Heap.new(nums)
        @heap.pop while @heap.size > k
    end

    def add(val)
        @heap.push(val)
        @heap.pop if @heap.size > @k
        @heap.peek
    end
end

This problem is the size-k template in class form. The constructor does the shrink once, and add repeats the push-then-trim cycle forever. The space stays at

O(k)

no matter how long the stream runs, which is the property that makes the design acceptable for live data.

Medium Problems

3. Kth Largest Element in an Array

LeetCode 215 | Difficulty: Medium

Brief: Return the kth largest element of an unsorted array.

Why this pattern: The same size-k min-heap works on a static array. Sort the whole array and you pay

O(N log N)

. The capped heap pays

O(N log k)

and never stores more than k values.

Key Insight: After every insert, pop the root if the heap holds more than k elements. The root is then the kth largest of everything processed, which is exactly the answer when the array is exhausted.

Complexity: Time

O(N log k)

, Space

O(k)

Visual:

    graph TD
    A["nums: 3, 2, 1, 5, 6, 4, k = 2"] --> B["Push 3, 2"]
    B --> C["Push 1, size > 2, pop 1"]
    C --> D["Push 5, pop 2"]
    D --> E["Push 6, pop 3"]
    E --> F["Push 4, pop 4"]
    F --> G["Root: 5, the 2nd largest"]

Code:

var findKthLargest = function(nums, k) {
    // min-heap capped at k keeps the k largest values.
    // the root is the kth largest of everything seen.
    const heap = new Heap();
    for (const num of nums) {
        heap.push(num);
        if (heap.size() > k) {
            heap.pop();
        }
    }
    return heap.peek();
};

import heapq

def find_kth_largest(nums, k):
    heap = []
    for num in nums:
        # keep the k largest values; the root of the
        # min-heap is the kth largest of everything seen
        heapq.heappush(heap, num)
        if len(heap) > k:
            heapq.heappop(heap)
    return heap[0]

class Solution {
    public int findKthLargest(int[] nums, int k) {
        // the default PriorityQueue is a min-heap
        PriorityQueue<Integer> heap = new PriorityQueue<>();
        for (int num : nums) {
            heap.offer(num);
            if (heap.size() > k) {
                heap.poll();
            }
        }
        return heap.peek();
    }
}

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        priority_queue<int, vector<int>, greater<int>> heap;
        for (int num : nums) {
            heap.push(num);
            if (heap.size() > k) {
                heap.pop();
            }
        }
        return heap.top();
    }
};

func findKthLargest(nums []int, k int) int {
    h := &IntMinHeap{}
    heap.Init(h)
    for _, num := range nums {
        heap.Push(h, num)
        if h.Len() > k {
            heap.Pop(h)
        }
    }
    return (*h)[0]
}

type IntMinHeap []int

func (h IntMinHeap) Len() int           { return len(h) }
func (h IntMinHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h IntMinHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *IntMinHeap) Push(x interface{}) {
    *h = append(*h, x.(int))
}

func (h *IntMinHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}

This is the stream version from the previous problem with the stream replaced by a loop. The guaranteed worst case is what distinguishes the heap here. Quickselect, the divide and conquer alternative covered on the divide and conquer page , has expected

O(N)

time but can degrade to

O(N^2)

with unlucky pivots. The heap never degrades.

4. Top K Frequent Elements

LeetCode 347 | Difficulty: Medium

Brief: Return the k most frequent elements from an array.

Why this pattern: Count frequencies first, then the top-K template applies with the frequency as the heap key instead of the raw value.

Key Insight: The heap stores (count, element) pairs and caps at k. The root is the least frequent element among the survivors, so it is the one that leaves when a more frequent element arrives.

Complexity: Time

O(N log k)

, Space

O(N)

Visual:

    graph TD
    A["nums: 1, 1, 1, 2, 2, 3"] --> B["freq: 1 has 3, 2 has 2, 3 has 1"]
    B --> C["k = 2"]
    C --> D["Push (3, 1), push (2, 2)"]
    D --> E["Push (1, 3), size > 2, pop (1, 3)"]
    E --> F["Return 1 and 2"]

Code:

var topKFrequent = function(nums, k) {
    const freq = new Map();
    for (const num of nums) {
        freq.set(num, (freq.get(num) || 0) + 1);
    }

    // pairs [count, num] in a min-heap; the root is the
    // least frequent element among the current k
    const heap = new Heap([], (a, b) => a[0] < b[0]);
    for (const [num, count] of freq) {
        heap.push([count, num]);
        if (heap.size() > k) {
            heap.pop();
        }
    }

    const result = [];
    while (heap.size() > 0) {
        result.push(heap.pop()[1]);
    }
    return result;
};

import heapq
from collections import Counter

def top_k_frequent(nums, k):
    freq = Counter(nums)
    heap = []
    for num, count in freq.items():
        # heap is keyed by count, so the root is the least
        # frequent element among the current k
        heapq.heappush(heap, (count, num))
        if len(heap) > k:
            heapq.heappop(heap)
    return [num for _, num in heap]

class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> freq = new HashMap<>();
        for (int num : nums) {
            freq.put(num, freq.getOrDefault(num, 0) + 1);
        }

        // int[] {count, num}, a min-heap by count
        PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> a[0] - b[0]);
        for (Map.Entry<Integer, Integer> e : freq.entrySet()) {
            heap.offer(new int[]{e.getValue(), e.getKey()});
            if (heap.size() > k) {
                heap.poll();
            }
        }

        int[] result = new int[k];
        for (int i = 0; i < k; i++) {
            result[i] = heap.poll()[1];
        }
        return result;
    }
}

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int, int> freq;
        for (int num : nums) {
            freq[num]++;
        }

        // pairs {count, num}, a min-heap by count
        priority_queue<pair<int, int>, vector<pair<int, int>>, greater<pair<int, int>>> heap;
        for (auto& e : freq) {
            heap.push({e.second, e.first});
            if (heap.size() > k) {
                heap.pop();
            }
        }

        vector<int> result;
        while (!heap.empty()) {
            result.push_back(heap.top().second);
            heap.pop();
        }
        return result;
    }
};

func topKFrequent(nums []int, k int) []int {
    freq := make(map[int]int)
    for _, num := range nums {
        freq[num]++
    }

    h := &FreqHeap{}
    heap.Init(h)
    for num, count := range freq {
        heap.Push(h, [2]int{count, num})
        if h.Len() > k {
            heap.Pop(h)
        }
    }

    result := make([]int, 0, k)
    for h.Len() > 0 {
        result = append(result, heap.Pop(h).([2]int)[1])
    }
    return result
}

type FreqHeap [][2]int

func (h FreqHeap) Len() int           { return len(h) }
func (h FreqHeap) Less(i, j int) bool { return h[i][0] < h[j][0] }
func (h FreqHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *FreqHeap) Push(x interface{}) {
    *h = append(*h, x.([2]int))
}

func (h *FreqHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}

The frequency map costs

O(N)

space, which is unavoidable because every distinct value appears in it. After the count pass, the heap work is the same as any top-K problem. The result order does not matter, since the problem accepts any ordering of the k answers.

5. K Closest Points to Origin

LeetCode 973 | Difficulty: Medium

Brief: Given an array of points, return the k closest points to the origin.

Why this pattern: k closest means k smallest distances, so the max-heap variation applies. The heap caps at k, and the root holds the largest distance among the survivors.

Key Insight: Compare by squared distance, not the raw distance. The square root is monotonic, so it cannot change the ordering, and skipping it avoids floating point entirely.

Complexity: Time

O(N log k)

, Space

O(k)

Visual:

    graph TD
    A["points: (1,3), (-2,2), (5,-1), (3,4)"] --> B["dists: 10, 8, 26, 25"]
    B --> C["k = 2"]
    C --> D["(5,-1) at distance 26 is popped"]
    D --> E["(3,4) at distance 25 is popped"]
    E --> F["Survivors: (1,3) and (-2,2)"]

Code:

var kClosest = function(points, k) {
    const dist = (p) => p[0] * p[0] + p[1] * p[1];

    // max-heap by distance keeps the k closest points. the
    // root is the farthest survivor and leaves when a
    // closer point arrives.
    const heap = new Heap([], (a, b) => dist(a) > dist(b));
    for (const p of points) {
        heap.push(p);
        if (heap.size() > k) {
            heap.pop();
        }
    }

    const result = [];
    while (heap.size() > 0) {
        result.push(heap.pop());
    }
    return result;
};

import heapq

def k_closest(points, k):
    heap = []
    for x, y in points:
        d = x * x + y * y
        # negate the distance so the root of the min-heap
        # is the farthest point among the current k
        heapq.heappush(heap, (-d, x, y))
        if len(heap) > k:
            heapq.heappop(heap)
    return [[x, y] for _, x, y in heap]

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        // max-heap by squared distance, so the root is the
        // farthest of the k survivors
        PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> dist(b) - dist(a));
        for (int[] p : points) {
            heap.offer(p);
            if (heap.size() > k) {
                heap.poll();
            }
        }
        return heap.toArray(new int[k][]);
    }

    private int dist(int[] p) {
        return p[0] * p[0] + p[1] * p[1];
    }
}

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
        auto dist = [](const vector<int>& p) {
            return p[0] * p[0] + p[1] * p[1];
        };
        // max-heap by squared distance
        auto cmp = [&](const vector<int>& a, const vector<int>& b) {
            return dist(a) < dist(b);
        };
        priority_queue<vector<int>, vector<vector<int>>, decltype(cmp)> heap(cmp);

        for (const auto& p : points) {
            heap.push(p);
            if (heap.size() > k) {
                heap.pop();
            }
        }

        vector<vector<int>> result;
        while (!heap.empty()) {
            result.push_back(heap.top());
            heap.pop();
        }
        return result;
    }
};

func kClosest(points [][]int, k int) [][]int {
    h := &PointHeap{}
    heap.Init(h)
    for _, p := range points {
        heap.Push(h, p)
        if h.Len() > k {
            heap.Pop(h)
        }
    }

    result := make([][]int, 0, k)
    for h.Len() > 0 {
        result = append(result, heap.Pop(h).([]int))
    }
    return result
}

func dist(p []int) int {
    return p[0]*p[0] + p[1]*p[1]
}

type PointHeap [][]int

func (h PointHeap) Len() int           { return len(h) }
func (h PointHeap) Less(i, j int) bool { return dist(h[i]) > dist(h[j]) }
func (h PointHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *PointHeap) Push(x interface{}) {
    *h = append(*h, x.([]int))
}

func (h *PointHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}

The squared distance trick is worth stating out loud in an interview. It removes the square root, keeps the comparison exact with integers, and changes nothing about the answer. The custom comparator is the only new piece compared to the plain top-K template.

Hard Problems

6. Merge k Sorted Lists

LeetCode 23 | Difficulty: Hard

Brief: Merge k sorted linked lists into one sorted linked list.

Why this pattern: The min-heap holds the current head of every list. The smallest head in the heap is always the next element of the merged result.

Key Insight: When a head is popped, the next node from that same list enters the heap. This keeps the heap at exactly k nodes and ensures the smallest available element is always at the root.

Complexity: Time

O(N log k)

, Space

O(k)

Visual:

    graph TD
    A["Heads 1, 1, 2 enter the min-heap"] --> B["Pop 1, push that list's next: 4"]
    B --> C["Pop 1, push that list's next: 3"]
    C --> D["Pop 2, push that list's next: 6"]
    D --> E["Heap holds 3, 4, 6"]
    E --> F{"Heap empty?"}
    F -->|No| B
    F -->|Yes| G["Result: 1, 1, 2, 3, 4, 4, 5, 6"]

Code:

var mergeKLists = function(lists) {
    class ListNode {
        constructor(val = 0, next = null) {
            this.val = val;
            this.next = next;
        }
    }

    // min-heap of list nodes ordered by value
    const heap = new Heap([], (a, b) => a.val < b.val);
    for (const head of lists) {
        if (head) {
            heap.push(head);
        }
    }

    const dummy = new ListNode();
    let current = dummy;
    while (heap.size() > 0) {
        current.next = heap.pop();
        current = current.next;
        if (current.next) {
            heap.push(current.next);
        }
    }
    return dummy.next;
};

import heapq

class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next

def merge_k_lists(lists):
    # the list index breaks ties between equal values,
    # because heap elements need a total order
    heap = []
    for i, head in enumerate(lists):
        if head:
            heapq.heappush(heap, (head.val, i, head))

    dummy = ListNode()
    current = dummy
    while heap:
        _, i, node = heapq.heappop(heap)
        current.next = node
        current = current.next
        if node.next:
            heapq.heappush(heap, (node.next.val, i, node.next))
    return dummy.next

class Solution {
    public ListNode mergeKLists(ListNode[] lists) {
        // min-heap of list nodes ordered by value
        PriorityQueue<ListNode> heap = new PriorityQueue<>((a, b) -> a.val - b.val);
        for (ListNode head : lists) {
            if (head != null) {
                heap.offer(head);
            }
        }

        ListNode dummy = new ListNode(0);
        ListNode current = dummy;
        while (!heap.isEmpty()) {
            current.next = heap.poll();
            current = current.next;
            if (current.next != null) {
                heap.offer(current.next);
            }
        }
        return dummy.next;
    }
}

class Solution {
public:
    ListNode* mergeKLists(vector<ListNode*>& lists) {
        // min-heap of list node pointers ordered by value
        auto cmp = [](ListNode* a, ListNode* b) { return a->val > b->val; };
        priority_queue<ListNode*, vector<ListNode*>, decltype(cmp)> heap(cmp);
        for (ListNode* head : lists) {
            if (head) {
                heap.push(head);
            }
        }

        ListNode dummy(0);
        ListNode* current = &dummy;
        while (!heap.empty()) {
            current->next = heap.top();
            heap.pop();
            current = current->next;
            if (current->next) {
                heap.push(current->next);
            }
        }
        return dummy.next;
    }
};

type ListNode struct {
    Val  int
    Next *ListNode
}

func mergeKLists(lists []*ListNode) *ListNode {
    h := &ListNodeHeap{}
    heap.Init(h)
    for _, head := range lists {
        if head != nil {
            heap.Push(h, head)
        }
    }

    dummy := &ListNode{}
    current := dummy
    for h.Len() > 0 {
        current.Next = heap.Pop(h).(*ListNode)
        current = current.Next
        if current.Next != nil {
            heap.Push(h, current.Next)
        }
    }
    return dummy.Next
}

type ListNodeHeap []*ListNode

func (h ListNodeHeap) Len() int           { return len(h) }
func (h ListNodeHeap) Less(i, j int) bool { return h[i].Val < h[j].Val }
func (h ListNodeHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *ListNodeHeap) Push(x interface{}) {
    *h = append(*h, x.(*ListNode))
}

func (h *ListNodeHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}

class ListNode
    attr_accessor :val, :next

    def initialize(val = 0, next_node = nil)
        @val = val
        @next = next_node
    end
end

def merge_k_lists(lists)
    # heap of [value, node] pairs, ordered by value only.
    # the node is carried along but never compared.
    heap = Heap.new { |a, b| a[0] < b[0] }
    lists.each do |head|
        heap.push([head.val, head]) if head
    end

    dummy = ListNode.new
    current = dummy
    while heap.size > 0
        current.next = heap.pop[1]
        current = current.next
        heap.push([current.next.val, current.next]) if current.next
    end
    dummy.next
end

The dummy node removes the special case of the empty result, the same trick as in linked list reversal problems. The pairwise merge alternative, which pairs lists and merges two at a time, appears on the divide and conquer page . Both are

O(N log k)

; the heap version keeps the k heads in one structure and is the answer most interviewers expect to hear.

7. Find Median from Data Stream

LeetCode 295 | Difficulty: Hard

Brief: Design a class that supports adding a number and finding the median of all numbers added so far.

Why this pattern: The two-heap split keeps the lower half in a max-heap and the upper half in a min-heap. The median is always at the boundary between the two roots.

Key Insight: Every insert pushes into the lower half, moves the largest lower value into the upper half, then moves it back if the sizes drift. After the dance, the lower half holds the extra element when the total count is odd, and the median is a root read either way.

Complexity: Time

O(log N)

per add,

O(1)

per median read, Space

O(N)

Visual:

    graph TD
    A["Add 1"] --> B["low: [1], high: []"]
    B --> C["Add 2"]
    C --> D["low: [1], high: [2]"]
    D --> E["Add 3"]
    E --> F["low: [1, 2], high: [3]"]
    F --> G["Median is the low root: 2"]

Code:

var MedianFinder = function() {
    // low is a max-heap, high is a min-heap
    this.low = new Heap([], (a, b) => a > b);
    this.high = new Heap();
};

MedianFinder.prototype.addNum = function(num) {
    // push into low, then move low's largest into high.
    // this keeps every low value below every high value.
    this.low.push(num);
    this.high.push(this.low.pop());
    // low holds the extra element when the count is odd
    if (this.low.size() < this.high.size()) {
        this.low.push(this.high.pop());
    }
};

MedianFinder.prototype.findMedian = function() {
    if (this.low.size() > this.high.size()) {
        return this.low.peek();
    }
    return (this.low.peek() + this.high.peek()) / 2;
};

import heapq

class MedianFinder:
    def __init__(self):
        # lower half as a max-heap (negated values), upper
        # half as a min-heap. every low value stays at or
        # below every high value.
        self.low = []
        self.high = []

    def add_num(self, num):
        # push into low, then move low's largest value into
        # high. this keeps the ordering invariant between
        # the two halves.
        heapq.heappush(self.low, -num)
        heapq.heappush(self.high, -heapq.heappop(self.low))
        # low holds the extra element when the count is odd
        if len(self.low) < len(self.high):
            heapq.heappush(self.low, -heapq.heappop(self.high))

    def find_median(self):
        if len(self.low) > len(self.high):
            return -self.low[0]
        return (-self.low[0] + self.high[0]) / 2

class MedianFinder {
    // lower half as a max-heap, upper half as a min-heap
    private PriorityQueue<Integer> low = new PriorityQueue<>(Collections.reverseOrder());
    private PriorityQueue<Integer> high = new PriorityQueue<>();

    public void addNum(int num) {
        low.offer(num);
        high.offer(low.poll());
        if (low.size() < high.size()) {
            low.offer(high.poll());
        }
    }

    public double findMedian() {
        if (low.size() > high.size()) {
            return low.peek();
        }
        return (low.peek() + high.peek()) / 2.0;
    }
}

class MedianFinder {
private:
    // lower half as a max-heap, upper half as a min-heap
    priority_queue<int> low;
    priority_queue<int, vector<int>, greater<int>> high;

public:
    void addNum(int num) {
        low.push(num);
        high.push(low.top());
        low.pop();
        if (low.size() < high.size()) {
            low.push(high.top());
            high.pop();
        }
    }

    double findMedian() {
        if (low.size() > high.size()) {
            return low.top();
        }
        return (low.top() + high.top()) / 2.0;
    }
};

type MedianFinder struct {
    low  *IntMaxHeap
    high *IntMinHeap
}

func Constructor() MedianFinder {
    return MedianFinder{
        low:  &IntMaxHeap{},
        high: &IntMinHeap{},
    }
}

func (mf *MedianFinder) AddNum(num int) {
    heap.Push(mf.low, num)
    heap.Push(mf.high, heap.Pop(mf.low).(int))
    if mf.low.Len() < mf.high.Len() {
        heap.Push(mf.low, heap.Pop(mf.high).(int))
    }
}

func (mf *MedianFinder) FindMedian() float64 {
    if mf.low.Len() > mf.high.Len() {
        return float64((*mf.low)[0])
    }
    return float64((*mf.low)[0]+(*mf.high)[0]) / 2.0
}

type IntMaxHeap []int

func (h IntMaxHeap) Len() int           { return len(h) }
func (h IntMaxHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h IntMaxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *IntMaxHeap) Push(x interface{}) {
    *h = append(*h, x.(int))
}

func (h *IntMaxHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}

type IntMinHeap []int

func (h IntMinHeap) Len() int           { return len(h) }
func (h IntMinHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h IntMinHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *IntMinHeap) Push(x interface{}) {
    *h = append(*h, x.(int))
}

func (h *IntMinHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}

class MedianFinder
    def initialize
        # lower half as a max-heap, upper half as a min-heap
        @low = Heap.new { |a, b| a > b }
        @high = Heap.new
    end

    def add_num(num)
        @low.push(num)
        @high.push(@low.pop)
        @low.push(@high.pop) if @low.size < @high.size
    end

    def find_median
        return @low.peek if @low.size > @high.size
        (@low.peek + @high.peek) / 2.0
    end
end

The two-step insert reads strangely at first, but the order matters. Pushing into low first and then moving low’s largest into high guarantees the ordering invariant, because low’s largest is always smaller than or equal to anything already in high. The rebalance step then repairs any size drift. Trace the insert with an even count and an odd count and the median read becomes obviously correct. Odd counts use the low root, and even counts average the two roots.

Hard problems are where patterns separate candidates who pass from candidates who get offers. Track your misses and review them on a spaced schedule so you do not repeat them in the real interview.

These seven problems cover the full range of heap applications. Start with the extraction loop in Last Stone Weight, internalize the size cap through the three top-K problems, and finish with the two-heap balance in Find Median from Data Stream. By the end, you should be able to look at a problem and say immediately which extreme needs to sit at the root, and which of the templates supplies it.

Done with these problems? The app has more, plus a spaced repetition system that brings problems back right before you would forget them. Continue your prep .