DSA Series: Heap (Blind 75, Java)

Overview

• Pattern cheat-sheet
• Primer: PriorityQueue in Java
• 1. Top K Frequent Elements (Medium · 347)
• 2. Find Median from Data Stream (Hard · 295)

Part of the DSA Blind 75 series. Track your progress on the Interview Prep page.

Targets: Visa, Mastercard, NPCI, Amex, PayU. Java 11+. All code compiled and tested against examples, edge cases and a randomized check (median).

Pattern cheat-sheet

Primer: PriorityQueue in Java

• java.util.PriorityQueue<E> is a binary min-heap by default (natural ordering). It is backed by an array; peek() is O(1), offer/add and poll are O(log n), remove(Object) and contains are O(n), constructing from a collection is O(n) (heapify).
• Max-heap: new PriorityQueue<>(Collections.reverseOrder()) or (a, b) -> Integer.compare(b, a).
• Custom comparator: new PriorityQueue<>((a, b) -> Integer.compare(a[1], b[1])) for int[]; Comparator.comparingInt(...), .thenComparing(...), .reversed() are cleaner for objects.
• Pitfalls
  • Never write (a, b) -> b - a on arbitrary ints: subtraction overflows (e.g. Integer.MIN_VALUE). Use Integer.compare / Long.compare.
  • Iterating a PriorityQueue (for, toString, stream) does not give sorted order; only repeated poll() does.
  • It is not thread-safe (use PriorityBlockingQueue) and does not allow null.
  • Ties are not stable; add a sequence number if FIFO among equals matters.
  • Mutating an element's priority after insertion breaks the heap invariant; remove and re-insert (O(n)) or use a lazy-deletion pattern (push a new entry, ignore stale ones when polled).
  • new PriorityQueue<>(n) takes an initial capacity, not a bound - it never evicts. To keep the top K you must poll() yourself when size() > k.
  • Top K largest uses a min-heap of size K (the root is the smallest of the K, i.e. the one to evict); top K smallest uses a max-heap of size K.
  • Comparators on Integer objects with == compare references; use .equals / compare.
  • Max-heap size-K heaps with peek() on empty queue return null (not an exception) - element()/remove() throw.
• Rule of thumb for "top K": sort O(n log n); size-K heap O(n log k); QuickSelect average O(n) (worst O(n^2)); bucket sort O(n) when keys are bounded.

1. Top K Frequent Elements (Medium · 347)

Problem statement. Given an integer array nums and an integer k, return the k most frequent elements (any order). The answer is guaranteed unique. Example: nums=[1,1,1,2,2,3], k=2 -> [1,2]. nums=[1], k=1 -> [1]. Constraints: 1 <= nums.length <= 10^5, -10^4 <= nums[i] <= 10^4, 1 <= k <= number of distinct elements. Follow-up: better than O(n log n).

How the interviewer asks / follow-ups.

• "Top K merchants by transaction count / top K most frequent error codes / top K BINs / hot keys / most-declined cards in the last hour", "trending products".
• Follow-ups: data arrives as a stream (maintain counts in a HashMap + heap with lazy updates, or a Count-Min Sketch + heap for approximate top-K at huge scale); data too large for one machine (shard by hash(key), compute local top-K, merge - note a global top-K needs each shard's full counts for its own keys, which hash partitioning guarantees); sliding window (time-bucketed counts); ties; return in sorted order; memory-bounded (Space-Saving / Misra-Gries algorithm); top K least frequent.
• Evaluates: frequency map first, then the heap size argument (K vs n), bucket sort for O(n), knowing QuickSelect, distributed thinking.

Approach 1: Brute force. Count frequencies, sort all distinct entries by frequency descending, take first k. Time O(n + d log d) where d = distinct count (<= n); space O(d).

Approach 2: Min-heap of size k. Key insight: you only need the k best, so keep a min-heap of size k keyed by frequency; when it grows past k, evict the smallest-frequency element (the root). At the end the heap holds the k most frequent. Dry run [1,1,1,2,2,3], k=2: freq {1:3, 2:2, 3:1}. Push 1 (size1). Push 2 (size2). Push 3 (size3 > 2) -> poll the min freq = 3(1). Remaining {1,2}.

Approach 3: Bucket sort by frequency (O(n)). Key insight: a frequency is between 1 and n, so make n+1 buckets indexed by frequency and walk from the highest bucket down collecting elements until k are gathered. Dry run same input: buckets[3]=[1], buckets[2]=[2], buckets[1]=[3]. Walk from index 6 down: skip empties; 3 -> take 1; 2 -> take 2; stop at k=2.

Approach 4: QuickSelect on the distinct-elements-by-frequency array: average O(d), worst O(d^2) (mitigate with random pivot). Mention only.

Java code

// Approach 2: min-heap of size k  -- O(n log k)
static int[] topKFrequent(int[] nums, int k) {
    Map<Integer, Integer> freq = new HashMap<>();
    for (int x : nums) freq.merge(x, 1, Integer::sum);
    PriorityQueue<Map.Entry<Integer, Integer>> pq =
        new PriorityQueue<>((a, b) -> Integer.compare(a.getValue(), b.getValue())); // min by frequency
    for (Map.Entry<Integer, Integer> e : freq.entrySet()) {
        pq.offer(e);
        if (pq.size() > k) pq.poll();            // evict the least frequent among the k+1
    }
    int[] res = new int[pq.size()];
    for (int i = res.length - 1; i >= 0; i--)    // poll gives ascending frequency; fill from the back
        res[i] = pq.poll().getKey();             // => result is most-frequent first
    return res;
}

// Approach 3: bucket sort  -- O(n)
@SuppressWarnings("unchecked")
static int[] topKBucket(int[] nums, int k) {
    Map<Integer, Integer> freq = new HashMap<>();
    for (int x : nums) freq.merge(x, 1, Integer::sum);
    List<Integer>[] buckets = new List[nums.length + 1];   // index = frequency
    for (Map.Entry<Integer, Integer> e : freq.entrySet()) {
        if (buckets[e.getValue()] == null) buckets[e.getValue()] = new ArrayList<>();
        buckets[e.getValue()].add(e.getKey());
    }
    int[] res = new int[Math.min(k, freq.size())];
    int idx = 0;
    for (int f = buckets.length - 1; f >= 1 && idx < res.length; f--) {
        if (buckets[f] == null) continue;
        for (int v : buckets[f]) {
            if (idx == res.length) break;
            res[idx++] = v;
        }
    }
    return res;
}

Complexity (n = array length, d = distinct values, d <= n)

• Brute force sort: Time O(n + d log d), Space O(d).
• Min-heap size k: Time O(n + d log k) - counting plus d heap operations of O(log k); Space O(d + k).
• Bucket sort: Time O(n) - counting, bucketing and one sweep; Space O(n).
• QuickSelect: average O(n), worst O(d^2).

Edge cases / pitfalls. k equals d (return everything); k larger than d (clamp, as in the bucket version); single element; negative numbers (don't use them as array indexes - use the map); a max-heap of all entries works too (O(d log d)) but is not the optimal; using a min-heap of size k keyed incorrectly (key by frequency, not by value); comparator subtraction overflow (frequencies are small here, but use Integer.compare anyway); ties at the k-th place are ruled out by the problem - ask what to do otherwise.

Pattern & related. Frequency map + top-K selection (heap / bucket / QuickSelect). Related: Kth Largest Element in an Array (215), Top K Frequent Words (692), K Closest Points to Origin (973), Sort Characters By Frequency (451), Kth Largest Element in a Stream (703), Merge K Sorted Lists (23).

2. Find Median from Data Stream (Hard · 295)

Problem statement. Design MedianFinder with addNum(int num) and findMedian(). The median is the middle value of the sorted list, or the mean of the two middle values if the size is even. Example: add 1, add 2 -> median 1.5; add 3 -> median 2.0. Constraints: -10^5 <= num <= 10^5, at most 5*10^4 calls, findMedian is only called when there is at least one element. Follow-ups: if all numbers are in [0,100]? If 99% are in [0,100]?

How the interviewer asks / follow-ups.

• "Continuously compute the median transaction amount (or median latency) over a stream of payments to flag outliers", "running median of settlement amounts per merchant", "p50 latency of an API".
• Follow-ups: values in a small range [0,100] -> counting array, O(1) add, O(range) median; 99% in the range -> counting array + buckets for under/overflow with two lists; sliding window median (480) - two heaps with lazy deletion or two TreeMaps/multisets; arbitrary percentiles (p99) -> approximate with t-digest / HdrHistogram / quantile sketches; distributed -> per-node sketches merged; removeNum; thread safety (lock or a single writer); data does not fit memory.
• Evaluates: the two-heap invariant, balancing logic, overflow in the average, knowing alternatives (sorted list with insertion O(n), balanced BST / order-statistic tree), and awareness of approximate quantile sketches in production.

Approach 1: Brute force. Store numbers in a list; on every findMedian sort and pick the middle. add O(1), median O(n log n). Better: keep the list sorted by binary-search insertion: add O(n) (shifting), median O(1).

Approach 2: Two heaps (optimal). Key insight: the median only depends on the largest element of the lower half and the smallest element of the upper half. Keep a max-heap lo (smaller half) and a min-heap hi (larger half) with sizes equal, or lo having exactly one more. Median = lo.peek() (odd) or the average of both tops (even). Insert algorithm that keeps both invariants (every element of lo <= every element of hi, and size(lo) - size(hi) in {0,1}): push the number into lo, move lo's max to hi (this guarantees ordering), then if hi became larger than lo, move hi's min back to lo. Dry run: add 1: lo=[1] -> move 1 to hi: lo=[], hi=[1] -> hi bigger, move back: lo=[1], hi=[]. add 2: lo=[1,2] -> move max 2 to hi: lo=[1], hi=[2]; sizes equal OK. median=(1+2)/2=1.5. add 3: lo=[1,3] -> move 3 to hi: lo=[1], hi=[2,3] -> hi bigger, move 2 back: lo=[1,2], hi=[3]. median=lo.peek()=2.

Java code

static class MedianFinder {
    private final PriorityQueue<Integer> lo = new PriorityQueue<>(Collections.reverseOrder()); // max-heap: lower half
    private final PriorityQueue<Integer> hi = new PriorityQueue<>();                           // min-heap: upper half

    public void addNum(int num) {
        lo.offer(num);                    // 1) put in lower half
        hi.offer(lo.poll());              // 2) its max goes to upper half => ordering invariant holds
        if (hi.size() > lo.size())        // 3) rebalance: lo keeps the extra element
            lo.offer(hi.poll());
    }

    public double findMedian() {
        if (lo.isEmpty()) throw new IllegalStateException("no data");
        if (lo.size() > hi.size()) return lo.peek();
        return ((long) lo.peek() + hi.peek()) / 2.0;   // long avoids int overflow of the sum
    }
}

Complexity (n = elements so far)

• Brute force (sort on query): add O(1), median O(n log n); space O(n).
• Sorted list with insertion: add O(n), median O(1); space O(n).
• Two heaps: add O(log n) - three heap operations; median O(1); space O(n).
• Counting array for bounded range R: add O(1), median O(R), space O(R).

Edge cases / pitfalls. Median on empty structure (define behavior - exception or NaN); integer overflow when adding two large ints before dividing (cast to long first, as above); integer division (/ 2 instead of / 2.0); duplicates are fine; first insert path; swapping the roles of heaps (max-heap must hold the lower half); forgetting that PriorityQueue.remove(Object) is O(n) when extending to sliding-window median; using size difference of 2 allowed (breaks peek logic).

Pattern & related. Two heaps (partition a multiset around a pivot position). Related: Sliding Window Median (480), IPO (502), Kth Largest Element in a Stream (703), Top K Frequent Elements (347), Merge K Sorted Lists (23), Meeting Rooms II (253).