Definition
- Max-heap (smaller half) – holds the lower half of all numbers, with the largest of them at the root for quick access.
- Min-heap (larger half) – holds the upper half of all numbers, with the smallest of them at the root.
- Balance invariant – the two heaps differ in size by at most one, so the boundary between them is always the median.
- Median query – when sizes are equal, the median is the average of the two roots; when one heap is larger, its root is the median.
- Insertion – each new number is placed in the correct heap based on the current max-heap root, then rebalanced if needed.
At the National Bank of India, transaction amounts arrive all day and the risk team wants the running median at any moment to spot drift. They split all amounts into two buckets: a low bucket (max-heap) holding the smaller half with its largest value on top, and a high bucket (min-heap) holding the larger half with its smallest value on top. The median sits right at the boundary between the two buckets.
Understanding the Concepts
What is the Median of a Number Stream problem?
Design a data structure that accepts numbers one at a time from a continuous stream and can return the median of all numbers seen so far at any point. The median of a sorted list is the middle element if the count is odd, or the average of the two middle elements if the count is even. Insertion and median query should both be efficient.
package com.bank.risk;
import java.util.Collections;
import java.util.PriorityQueue;
public class MedianFinder {
private PriorityQueue<Integer> maxHeap; // smaller half
private PriorityQueue<Integer> minHeap; // larger half
public MedianFinder() {
maxHeap = new PriorityQueue<>(Collections.reverseOrder());
minHeap = new PriorityQueue<>();
}
public void addNum(int amount) {
if (maxHeap.isEmpty() || amount <= maxHeap.peek()) {
maxHeap.offer(amount);
} else {
minHeap.offer(amount);
}
rebalance();
}
private void rebalance() {
if (maxHeap.size() > minHeap.size() + 1) {
minHeap.offer(maxHeap.poll());
} else if (minHeap.size() > maxHeap.size()) {
maxHeap.offer(minHeap.poll());
}
}
public double findMedian() {
if (maxHeap.size() == minHeap.size()) {
return (maxHeap.peek() + minHeap.peek()) / 2.0;
}
return maxHeap.peek();
}
}
Why use two heaps instead of sorting?
Re-sorting the entire stream on every query costs O(n log n) per query, which is impractical for a live stream. With two heaps, each insertion costs O(log n) for the heap operation and each median query costs O(1) by peeking at the tops. For a stream of one million numbers, the two-heap approach does about 20 million operations total, while re-sorting would do about 20 trillion.
// Comparison:
// Re-sort per query: O(n log n) per query, O(n^2 log n) total for n queries
// Two heaps: O(log n) per insert, O(1) per query, O(n log n) total
// For n=1,000,000: two heaps is about 1,000,000x faster
How does insertion decide which heap receives the number?
If the max-heap is empty or the new number is at most the max-heap's root, the number belongs in the smaller half, so push it to the max-heap. Otherwise, it belongs in the larger half, so push it to the min-heap. This ensures that every element in the max-heap is less than or equal to every element in the min-heap.
package com.bank.risk;
public class InsertionDemo {
public static void trace() {
// Stream: 500, 200, 800, 300
// Step 1: 500 -> maxHeap empty, push to maxHeap
// maxHeap=[500], minHeap=[]
// Step 2: 200 <= 500, push to maxHeap
// maxHeap=[500,200], minHeap=[], rebalance -> move 500
// maxHeap=[200], minHeap=[500]
// Step 3: 800 > 200, push to minHeap
// maxHeap=[200], minHeap=[500,800]
// rebalance -> move 500 to maxHeap
// maxHeap=[500,200], minHeap=[800]
// Step 4: 300 <= 500, push to maxHeap
// maxHeap=[500,300,200], minHeap=[800]
// rebalance -> move 500
// maxHeap=[300,200], minHeap=[500,800]
// Median = (300+500)/2 = 400
}
}
How does rebalancing maintain the size invariant?
After each insertion, check the sizes. If the max-heap has grown more than one larger than the min-heap, move its root to the min-heap. If the min-heap has grown larger than the max-heap, move its root to the max-heap. This convention keeps the max-heap either equal in size to the min-heap or exactly one larger, so the median is always accessible from the tops.
package com.bank.risk;
public class RebalanceDemo {
public static void showRebalance() {
// Before: maxHeap=[500, 200], minHeap=[800]
// maxHeap.size()=2 > minHeap.size()+1=2? No
// minHeap.size()=1 > maxHeap.size()=2? No
// No rebalance needed
// Before: maxHeap=[500, 300, 200], minHeap=[800]
// maxHeap.size()=3 > minHeap.size()+1=2? Yes
// Move maxHeap root (500) to minHeap
// After: maxHeap=[300, 200], minHeap=[500, 800]
}
}
How do you read the median?
If both heaps are the same size, the stream has an even number of elements, so the median is the average of the max-heap root (largest of the smaller half) and the min-heap root (smallest of the larger half). If the max-heap is larger, it holds the extra middle element, so its root is the median.
package com.bank.risk;
public class MedianReadDemo {
public static void main(String[] args) {
MedianFinder mf = new MedianFinder();
mf.addNum(500);
System.out.println("Median after 500: " + mf.findMedian());
mf.addNum(200);
System.out.println("Median after 200: " + mf.findMedian());
mf.addNum(800);
System.out.println("Median after 800: " + mf.findMedian());
mf.addNum(300);
System.out.println("Median after 300: " + mf.findMedian());
}
}
Median after 500: 500.0
Median after 200: 350.0
Median after 800: 500.0
Median after 300: 400.0
What about using BigDecimal for banking amounts?
In a real banking system, amounts are BigDecimal to avoid floating-point precision loss. The max-heap needs a reverse-order comparator for BigDecimal, and the even-count median must be computed with BigDecimal division at a controlled scale and rounding mode.
package com.bank.risk;
import java.math.BigDecimal;
import java.math.RoundingMode;
import java.util.Collections;
import java.util.PriorityQueue;
public class BigDecimalMedianFinder {
private PriorityQueue<BigDecimal> maxHeap =
new PriorityQueue<>(Collections.reverseOrder());
private PriorityQueue<BigDecimal> minHeap =
new PriorityQueue<>();
public void addAmount(BigDecimal amount) {
if (maxHeap.isEmpty() || amount.compareTo(maxHeap.peek()) <= 0) {
maxHeap.offer(amount);
} else {
minHeap.offer(amount);
}
if (maxHeap.size() > minHeap.size() + 1) minHeap.offer(maxHeap.poll());
else if (minHeap.size() > maxHeap.size()) maxHeap.offer(minHeap.poll());
}
public BigDecimal findMedian() {
if (maxHeap.size() == minHeap.size()) {
return maxHeap.peek().add(minHeap.peek())
.divide(BigDecimal.valueOf(2), 2, RoundingMode.HALF_UP);
}
return maxHeap.peek();
}
}
What is the time and space complexity?
| Operation | Time | Space |
|---|---|---|
| addNum | O(log n) | O(1) per call |
| findMedian | O(1) | O(1) |
| Total for n inserts | O(n log n) | O(n) for both heaps |
| Re-sort alternative | O(n^2 log n) total | O(n) |
What are common mistakes?
The most common mistake is choosing the wrong rebalance direction, which causes the max-heap to grow too large or the min-heap to become larger than the max-heap. Another mistake is using floating-point arithmetic for money instead of BigDecimal.
package com.bank.risk;
public class CommonMistakes {
// WRONG: rebalance allows minHeap to be 2 larger
public static void wrongRebalance() {
// Only checking maxHeap > minHeap + 1
// but NOT checking minHeap > maxHeap
// Fix: add the second check
}
// WRONG: using double for money
public static void wrongMoney() {
double a = 0.1 + 0.2; // 0.30000000000000004, not 0.3!
// Fix: use BigDecimal for all monetary calculations
}
}
What are common interview follow-up questions?
Interviewers may ask how to handle removal of elements from the stream, how to find the median in a sliding window of size k (which requires a more complex augmented structure), and whether the approach works for a distributed system where numbers arrive on different machines.
// Follow-up: sliding window median
// - Use two TreeMaps (multisets) instead of heaps
// - TreeMap supports O(log n) removal by value
// - Maintain the same size invariant
// - On each window slide: add new element, remove expired element, rebalance
Visual Overview
Two-Heap Structure
==================
┌─────────────┐ ┌─────────────┐
│ Max-Heap │ │ Min-Heap │
│ (smaller ½) │ │ (larger ½) │
│ │ │ │
│ ┌───┐ │ │ ┌───┐ │
│ │300│ root │ │ │500│ root│
│ └─┬─┘ │ │ └─┬─┘ │
│ ┌──┴──┐ │ │ ┌──┴──┐ │
│ │200 │ │ │ │800 │ │
│ └─────┘ │ │ └─────┘ │
└─────────────┘ └─────────────┘
▲ ▲
│ │
Largest of Smallest of
smaller half larger half
│ │
└────── median ──────┘
(300 + 500) / 2 = 400
Insertion Flow
==============
New number arrives
│
▼
┌──────────────────┐
│ num <= maxHeap │
│ root? │
└───┬──────┬───────┘
Yes No
│ │
▼ ▼
Push to Push to
maxHeap minHeap
│ │
└───┬───┘
▼
┌──────────────────┐
│ Rebalance │
│ |size diff| <= 1 │
└──────────────────┘
│
▼
Median ready
Key Takeaway
Finding the median of a number stream is the textbook two-heap problem. You keep a max-heap that holds the smaller half of all numbers and a min-heap that holds the larger half, and you keep their sizes within one of each other. The median then lives at the boundary: when the counts are equal it is the average of the two heap tops, and when they differ it is the top of the larger heap. Each insertion is a logarithmic heap operation and each median query is a constant-time peek.
The insertion rule is straightforward: if the new number is at most the max-heap root, it belongs in the smaller half; otherwise it belongs in the larger half. After each insertion you rebalance by moving a root across if one heap grew two larger than the other. This ensures the boundary always sits at the true median position.
At the National Bank of India, the risk team used this to maintain the running median transaction amount over a live stream, watching for drift in typical transaction sizes in real time. They stressed that amounts must be BigDecimal, so the max-heap used a reverse-order comparator and the even-count median was computed with BigDecimal division at a fixed scale and rounding mode, never floating-point, because money must stay exact.
The key advantage over re-sorting is performance: two heaps give O(n log n) total for n insertions with O(1) per query, while re-sorting gives O(n^2 log n) total. For a live stream of millions of transactions, that difference is the line between real-time and unusable.
Remember one sentence: smaller half in a max-heap, larger half in a min-heap, balanced within one, and the median lives at the boundary between the two roots.
---