Find Median From Data Stream

The problem

Keep track of the median of a growing collection of numbers. operations lists the calls to carry out, in order, and values holds the arguments of each call at the same position. "addNum" [value] adds one number; the same number may be added several times, and every copy counts. "findMedian" [] asks for the median of all the numbers added so far: the middle one when they are put in order, or the average of the two middle ones when their count is even. Return a list with one result per call: the median for "findMedian" and null for every other call, including the opening "MedianFinder".

What to write

function solve(operations: string[], values: number[][]): (number | null)[]

Input
operations = ["MedianFinder","addNum","addNum","findMedian","addNum","findMedian"]
values = [[],[1],[2],[],[3],[]]

Output
[null,null,null,1.5,null,2]

Input
operations = ["MedianFinder","addNum","findMedian","addNum","findMedian"]
values = [[],[-1],[],[-2],[]]

Output
[null,null,-1,null,-1.5]

Constraints

The median query is called only after at least one value is inserted.

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Use the constraints

   Numbers keep arriving and the median is asked for again and again, so sorting everything from scratch for each "findMedian" repeats most of the work. Notice what the median needs: only the one or two numbers in the middle, never the order of the others.

2. Hint 2: Ask a smaller question

   The numbers so far, in order, are 1, 3, 5, 7, 9, so the median is 5. Now 6 arrives. Which of the old numbers do you need to look at to find the new median? Would the answer change if you did not know whether 1 or 3 came first?

3. Hint 3: Name the pattern

   Split the numbers into a smaller half and a larger half. The median only needs the largest of the smaller half and the smallest of the larger half. A heap is a collection that gives its largest or its smallest item at once, so keep one heap per half. This is the two heaps pattern. Read the lesson on two heaps.

4. Hint 4: State the invariant

   Two things stay true after every "addNum". Every number in the smaller half is less than or equal to every number in the larger half. And the smaller half holds as many numbers as the larger half, or exactly one more.

5. Hint 5: Sketch the pseudocode

   Keep two heaps: small, whose top is its largest number (a max-heap), and large, whose top is its smallest (a min-heap). addNum: push the number into small; move the top of small into large; if large now holds more, move its top into small. findMedian: the top of small if small holds more, otherwise the average of both tops.

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Approach

Keep the smaller half of the numbers in a max-heap, whose top is its largest number, and the larger half in a min-heap, whose top is its smallest. The smaller half stays the same size or one number bigger. The median is the top of the smaller half, or the average of both tops when the sizes are equal.

The slow way

Keep the numbers in a list and sort it for every "findMedian", then read the middle. It is correct, but each call sorts all the numbers again, about n log n steps for n numbers. Keeping the list sorted and inserting each new number in its place avoids the sort, yet one insert can shift up to n numbers.

The insight

The median only needs the middle of the sorted order, never the whole order. Keep the numbers as a smaller half and a larger half. A heap is a collection that hands you its largest number (a max-heap) or its smallest (a min-heap) at once, and adds or removes one in about log n steps. One heap per half keeps both middle numbers at hand.

Step by step

1. Keep the smaller half and the larger half of the numbers in two heaps. The median sits where the two halves meet. So only the largest number of the smaller half and the smallest of the larger half are ever read, and a heap keeps exactly that one number on top.
2. Carry out the operations in order and collect one result for each. The calls come as two lists, the names and their arguments. "addNum" changes the heaps and "findMedian" reads them, and the answer needs one entry per call, null when the call returns nothing.
3. A new number first goes into the smaller half. You do not know yet which half the number belongs to. Putting it in the smaller half and letting the next step correct that means no comparison has to be written here.
4. Move the largest of the smaller half over to the larger half. The new number may be bigger than some numbers of the larger half. Sending the largest of the smaller half across fixes that: afterwards no number in the smaller half is bigger than any number in the larger half.
5. If the larger half now holds more numbers, move its smallest back. The move made the larger half one number bigger. If it now holds more than the smaller half, returning its smallest number keeps the smaller half equal in size or one ahead.
6. Odd count: the median is the largest of the smaller half. With an odd count the smaller half holds one number more than the larger half, so its largest number is the middle one of them all.
7. Even count: the median is the average of the two middle numbers. With an even count the halves are the same size, so the two middle numbers are the largest of the smaller half and the smallest of the larger half.

The same numbered steps are the comments in each solution.

function solve(operations, values) {
  // JavaScript has no built-in heap. These two helpers keep an array as a
  // min-heap: the smallest value is at index 0, and the value at index i is
  // never larger than its two children, at 2 * i + 1 and 2 * i + 2.
  function heapPush(heap, value) {
    // Add the value at the end, then swap it upward past any larger parent.
    heap.push(value);
    let i = heap.length - 1;
    while (i > 0) {
      const parent = Math.floor((i - 1) / 2);
      if (heap[parent] <= heap[i]) break;
      [heap[parent], heap[i]] = [heap[i], heap[parent]];
      i = parent;
    }
  }
  function heapPop(heap) {
    // Remove the top, put the last value in its place, then swap that value
    // downward with its smaller child until no child is smaller.
    const top = heap[0];
    const last = heap.pop();
    if (heap.length === 0) return top;
    heap[0] = last;
    let i = 0;
    while (true) {
      let child = 2 * i + 1;
      if (child + 1 < heap.length && heap[child + 1] < heap[child]) child++;
      if (child >= heap.length || heap[i] <= heap[child]) break;
      [heap[i], heap[child]] = [heap[child], heap[i]];
      i = child;
    }
    return top;
  }
  // 1. Keep the smaller half and the larger half of the numbers in two heaps
  // (The smaller half is stored negated, so the top of its min-heap is the
  // largest number of that half.)
  const small = [];
  const large = [];
  // 2. Carry out the operations in order and collect one result for each
  return operations.map((op, i) => {
    // "MedianFinder" only creates the two heaps, so it has no result.
    if (op === "MedianFinder") return null;
    if (op === "addNum") {
      // 3. A new number first goes into the smaller half
      heapPush(small, -values[i][0]);
      // 4. Move the largest of the smaller half over to the larger half
      heapPush(large, -heapPop(small));
      // 5. If the larger half now holds more numbers, move its smallest back
      if (large.length > small.length) heapPush(small, -heapPop(large));
      return null;
    }
    // 6. Odd count: the median is the largest of the smaller half
    if (small.length > large.length) return -small[0];
    // 7. Even count: the median is the average of the two middle numbers
    return (-small[0] + large[0]) / 2;
  });
}

function solve(operations: string[], values: number[][]): (number | null)[] {
  // TypeScript has no built-in heap. These two helpers keep an array as a
  // min-heap: the smallest value is at index 0, and the value at index i is
  // never larger than its two children, at 2 * i + 1 and 2 * i + 2.
  function heapPush(heap: number[], value: number): void {
    // Add the value at the end, then swap it upward past any larger parent.
    heap.push(value);
    let i = heap.length - 1;
    while (i > 0) {
      const parent = Math.floor((i - 1) / 2);
      if (heap[parent] <= heap[i]) break;
      [heap[parent], heap[i]] = [heap[i], heap[parent]];
      i = parent;
    }
  }
  function heapPop(heap: number[]): number {
    // Remove the top, put the last value in its place, then swap that value
    // downward with its smaller child until no child is smaller.
    const top = heap[0];
    const last = heap.pop()!;
    if (heap.length === 0) return top;
    heap[0] = last;
    let i = 0;
    while (true) {
      let child = 2 * i + 1;
      if (child + 1 < heap.length && heap[child + 1] < heap[child]) child++;
      if (child >= heap.length || heap[i] <= heap[child]) break;
      [heap[i], heap[child]] = [heap[child], heap[i]];
      i = child;
    }
    return top;
  }
  // 1. Keep the smaller half and the larger half of the numbers in two heaps
  // (The smaller half is stored negated, so the top of its min-heap is the
  // largest number of that half.)
  const small: number[] = [];
  const large: number[] = [];
  // 2. Carry out the operations in order and collect one result for each
  return operations.map((op, i) => {
    // "MedianFinder" only creates the two heaps, so it has no result.
    if (op === "MedianFinder") return null;
    if (op === "addNum") {
      // 3. A new number first goes into the smaller half
      heapPush(small, -values[i][0]);
      // 4. Move the largest of the smaller half over to the larger half
      heapPush(large, -heapPop(small));
      // 5. If the larger half now holds more numbers, move its smallest back
      if (large.length > small.length) heapPush(small, -heapPop(large));
      return null;
    }
    // 6. Odd count: the median is the largest of the smaller half
    if (small.length > large.length) return -small[0];
    // 7. Even count: the median is the average of the two middle numbers
    return (-small[0] + large[0]) / 2;
  });
}

import heapq


def solve(operations, values):
    # 1. Keep the smaller half and the larger half of the numbers in two heaps
    # (heapq is a min-heap. The smaller half is stored negated, so the top of
    # its heap is the largest number of that half.)
    small = []
    large = []
    # 2. Carry out the operations in order and collect one result for each
    out = []
    for op, args in zip(operations, values):
        result = None
        if op == "addNum":
            # 3. A new number first goes into the smaller half
            heapq.heappush(small, -args[0])
            # 4. Move the largest of the smaller half over to the larger half
            heapq.heappush(large, -heapq.heappop(small))
            # 5. If the larger half now holds more numbers, move its smallest back
            if len(large) > len(small):
                heapq.heappush(small, -heapq.heappop(large))
        elif op == "findMedian":
            # 6. Odd count: the median is the largest of the smaller half
            if len(small) > len(large):
                result = -small[0]
            # 7. Even count: the median is the average of the two middle numbers
            else:
                result = (-small[0] + large[0]) / 2
        # "MedianFinder" and "addNum" leave the result as None.
        out.append(result)
    return out

Example 1, traced

Input
operations = ["MedianFinder","addNum","addNum","findMedian","addNum","findMedian"]
values = [[],[1],[2],[],[3],[]]

Output
[null,null,null,1.5,null,2]

Why it works

After every "addNum" two things hold. No number in the smaller half is bigger than a number in the larger half, because the largest of the smaller half always crosses over. And the smaller half is the same size as the larger half or one ahead, because of the size check. So the middle of the sorted numbers sits at the two tops.

Time and space

Time O(log n) insertion; O(1) median, extra space O(n). An "addNum" makes at most five pushes and pops. Each moves a number along one path of a heap, about log n steps for n numbers, so an insertion is O(log n). "findMedian" only reads the two tops, in constant time. Each number is stored once: O(n) space.

n and m are input sizes, h is a tree's height, V and E are a graph's vertices and edges. The bounds describe this reference solution.

Watch out

• Do not put a new number straight into the half that has fewer numbers. In example 1 the 3 would then join the 1 in the smaller half, and the last median comes out as 3 instead of 2.
• The references keep the smaller half with flipped signs, because their heap gives its smallest value first. Flip the sign back whenever a number leaves that heap. Forget it on the move across and example 1 answers -1.5, not 1.5.

Pattern

Two heaps. Builds on Kth Largest Element In a Stream.

Keep going