Skip to content
DSA heaps 12 min read

Heap Exercises: Answer Sheet & Worked Solutions

This is the answer sheet for the heap exercises. Each solution gives the key idea, the optimal complexity, the code in all four languages (use the switcher), and a common pitfall. Try the problems first — then check your work here.

JavaScript note: JS has no built-in heap, so each JS solution includes a small comparator-based Heap class (cmp(a, b) < 0 means a has higher priority). It’s the same sift-up/sift-down logic from heap operations. C++ uses std::priority_queue, Java uses PriorityQueue, and Python uses heapq — see priority queues.

Q1. Kth largest element in an array

Idea: Keep a min-heap of size k. Push every element; whenever the heap exceeds k, pop the smallest. The heap ends holding the k largest values, and its root is the kth largest.

Complexity: O(n log k) time, O(k) space.

#include <queue>
#include <vector>
int findKthLargest(std::vector<int>& nums, int k) {
    std::priority_queue<int, std::vector<int>, std::greater<int>> minHeap;
    for (int x : nums) {
        minHeap.push(x);
        if ((int)minHeap.size() > k) minHeap.pop();
    }
    return minHeap.top();
}

Pitfall: “kth largest” means kth in sorted order, not the kth distinct value — duplicates count. Also, a max-heap of all n elements then k pops is O(n + k log n); the size-k min-heap is cheaper in space and scales to streams.

Q2. Last stone weight

Idea: A max-heap always hands you the two heaviest stones. Smash them, push back the difference if nonzero, and repeat until one or zero stones remain.

Complexity: O(n log n) time, O(n) space.

#include <queue>
#include <vector>
int lastStoneWeight(std::vector<int>& stones) {
    std::priority_queue<int> pq(stones.begin(), stones.end()); // max-heap
    while (pq.size() > 1) {
        int a = pq.top(); pq.pop();
        int b = pq.top(); pq.pop();
        if (a != b) pq.push(a - b);
    }
    return pq.empty() ? 0 : pq.top();
}

Pitfall: Forgetting the empty-heap case returns garbage — return 0 when no stones remain. With Python’s min-only heapq, you must negate on the way in and out.

Q3. Top K frequent elements

Idea: Count frequencies with a hash map, then keep a min-heap of size k keyed by count. Elements with the highest counts survive; pop whenever the heap exceeds k.

Complexity: O(m log k) time where m is the number of distinct values, O(m) space.

#include <queue>
#include <vector>
#include <unordered_map>
std::vector<int> topKFrequent(std::vector<int>& nums, int k) {
    std::unordered_map<int,int> freq;
    for (int x : nums) freq[x]++;
    using P = std::pair<int,int>;                 // {count, value}
    std::priority_queue<P, std::vector<P>, std::greater<P>> pq; // min by count
    for (auto& [val, cnt] : freq) {
        pq.push({cnt, val});
        if ((int)pq.size() > k) pq.pop();
    }
    std::vector<int> res;
    while (!pq.empty()) { res.push_back(pq.top().second); pq.pop(); }
    return res;
}

Pitfall: Sorting all distinct values is O(m log m); the size-k heap is O(m log k). If k is close to m, bucket sort by count is an even cleaner O(n). Python’s heapq.nlargest(k, freq, key=freq.get) is the idiomatic one-liner.

Q4. K closest points to origin

Idea: Keep a max-heap of size k keyed by squared distance. When the heap exceeds k, pop the farthest. The k closest survive. Compare squared distances — no sqrt needed.

Complexity: O(n log k) time, O(k) space.

#include <queue>
#include <vector>
std::vector<std::vector<int>> kClosest(std::vector<std::vector<int>>& points, int k) {
    std::priority_queue<std::pair<int,int>> pq; // max-heap of {dist2, index}
    for (int i = 0; i < (int)points.size(); i++) {
        int d = points[i][0]*points[i][0] + points[i][1]*points[i][1];
        pq.push({d, i});
        if ((int)pq.size() > k) pq.pop();
    }
    std::vector<std::vector<int>> res;
    while (!pq.empty()) { res.push_back(points[pq.top().second]); pq.pop(); }
    return res;
}

Pitfall: Taking the square root is wasteful and risks float rounding — ordering by x*x + y*y gives identical results. Use a max-heap (not min) so the size-k version can evict the farthest point.

Q5. Merge k sorted lists

Idea: Seed a min-heap with the first element of each list. Repeatedly pop the smallest, append it, and push the next element from the same list. The heap never holds more than k entries.

Complexity: O(N log k) time (N = total elements), O(k) space for the heap.

#include <queue>
#include <vector>
#include <tuple>
std::vector<int> mergeKSorted(std::vector<std::vector<int>>& lists) {
    using T = std::tuple<int,int,int>; // {value, listIdx, elemIdx}
    std::priority_queue<T, std::vector<T>, std::greater<T>> pq;
    for (int i = 0; i < (int)lists.size(); i++)
        if (!lists[i].empty()) pq.push({lists[i][0], i, 0});
    std::vector<int> res;
    while (!pq.empty()) {
        auto [val, li, ei] = pq.top(); pq.pop();
        res.push_back(val);
        if (ei + 1 < (int)lists[li].size())
            pq.push({lists[li][ei + 1], li, ei + 1});
    }
    return res;
}

Pitfall: Pushing entire lists into the heap makes it O(N log N) and defeats the purpose — push only the current front of each list. Skip empty lists when seeding, and include the list/element indices so you know where to pull the next value.

Q6. Find median from a data stream

Idea: Keep two heaps split at the median: a max-heap lo for the lower half and a min-heap hi for the upper half. Keep them balanced so lo has the same size as hi or one more. The median is lo’s top (odd count) or the average of both tops (even count).

Complexity: O(log n) per addNum, O(1) per findMedian, O(n) space.

#include <queue>
#include <vector>
class MedianFinder {
    std::priority_queue<int> lo;                                   // max-heap: lower half
    std::priority_queue<int, std::vector<int>, std::greater<int>> hi; // min-heap: upper half
public:
    void addNum(int x) {
        lo.push(x);
        hi.push(lo.top()); lo.pop();           // shuttle the new top across the split
        if (hi.size() > lo.size()) { lo.push(hi.top()); hi.pop(); }
    }
    double findMedian() {
        if (lo.size() > hi.size()) return lo.top();
        return (lo.top() + hi.top()) / 2.0;
    }
};

Pitfall: Push-to-lo then move lo’s top to hi guarantees the two halves stay correctly ordered relative to each other; skipping that step lets a small number sit in hi. For an even count, average both tops as a float — integer division truncates the median.

That’s the full answer sheet. Back to the heap exercises, revisit priority queues, or continue to tries.

Last updated June 25, 2026
Was this helpful?