Find Median From Data Stream

NeetCode: Find Median From Data Stream · LeetCode: 295. Find Median From Data Stream · Difficulty: Hard

Design a structure MedianFinder with addNum(num) and findMedian(). The median is the middle value of the sorted numbers, or the mean of the two middle values when the count is even.

Example

Text

addNum(1); addNum(2)
findMedian() -> 1.5
addNum(3)
findMedian() -> 2.0

Brute force: sort at every query

Java

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

class MedianFinder {
    private final List<Integer> nums = new ArrayList<>();

    public MedianFinder() {
    }

    public void addNum(int num) {
        nums.add(num);
    }

    public double findMedian() {
        List<Integer> sorted = new ArrayList<>(nums);
        Collections.sort(sorted);
        int n = sorted.size();
        if (n % 2 == 1) {
            return sorted.get(n / 2);
        }
        return ((double) sorted.get(n / 2 - 1) + sorted.get(n / 2)) / 2.0;
    }
}

public class Main {
    public static void main(String[] args) {
        MedianFinder m = new MedianFinder();
        m.addNum(1);
        m.addNum(2);
        System.out.println(m.findMedian());
        m.addNum(3);
        System.out.println(m.findMedian());
    }
}

Output

1.5
2.0

The sum of the two middle values is computed in double, because adding two int values near Integer.MAX_VALUE would overflow. findMedian returns a double, so the whole-number median prints as 2.0.

Time: addNum is O(1), findMedian is O(n log n) (it sorts a copy). Space: O(n).

Optimal: two heaps

Split the numbers into a lower half in a max-heap and an upper half in a min-heap. Java's PriorityQueue is a min-heap by default, so the lower half uses Collections.reverseOrder(). Keep them balanced so the lower half has the same size as the upper half, or one more. Then the median is the top of the lower heap (odd count) or the average of the two tops (even count).

To insert, push into the lower heap, move that heap's largest into the upper heap (which guarantees everything below is <= everything above), and if the upper heap is now bigger, move its smallest back down.

Java

import java.util.Collections;
import java.util.PriorityQueue;

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

    public MedianFinder() {
    }

    public void addNum(int num) {
        lower.add(num);
        upper.add(lower.poll());
        if (upper.size() > lower.size()) {
            lower.add(upper.poll());
        }
    }

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

public class Main {
    public static void main(String[] args) {
        MedianFinder m = new MedianFinder();
        m.addNum(1);
        m.addNum(2);
        System.out.println(m.findMedian());
        m.addNum(3);
        System.out.println(m.findMedian());
        m.addNum(10);
        System.out.println(m.findMedian());
    }
}

Output

1.5
2.0
2.5

peek returns the top element without removing it, and poll removes and returns it. Time: addNum is O(log n), findMedian is O(1). Space: O(n).