Find median in a row wise sorted matrix

The problem, simply

You get a grid with m rows and n columns. Each row is sorted from small to big. The columns are not sorted. The total count m * n is odd.

Give back the median: the number that would sit exactly in the middle if you poured all the numbers into one line and sorted it. With m * n numbers, that is the (m * n + 1) / 2th smallest.

A story to remember it

Picture it

Examples

Way 1: the simple way

Pour every number into one long line, sort it, and pick the middle.

1. Copy all m * n numbers into an array.
2. Sort the array.
3. Return the number at index (m * n) / 2.

import java.util.*;

class Solution {
    public int findMedian(int[][] matrix) {
        int m = matrix.length, n = matrix[0].length;
        int[] all = new int[m * n];
        int k = 0;
        for (int[] row : matrix)
            for (int value : row) all[k++] = value;   // pour everything into one line

        Arrays.sort(all);
        return all[(m * n) / 2];                      // the middle seat
    }
}

How long it takes: sorting m * n numbers is O(m*n*log(m*n)). Extra room: a copy of the whole grid, O(m*n). It ignores that the rows were already sorted.

Way 2: the clever way

Do not search for where the median sits. Search for what value it has.

The median is the smallest number x such that at least need = (m * n + 1) / 2 numbers in the grid are x or less. If we guess some mid:

• If at least need numbers are mid or less, the median is mid or smaller.
• If fewer than need numbers are, the median is bigger than mid.

Counting is fast, because each row is sorted: find where the numbers bigger than mid start in each row with a halving search (an upper bound).

1. lo = the smallest first number of any row. hi = the biggest last number of any row.
2. While lo < hi: let mid be halfway. Count numbers <= mid in all rows.
3. If the count is less than need, set lo = mid + 1. Otherwise set hi = mid.
4. When lo == hi, that is the median.

class Solution {
    public int findMedian(int[][] matrix) {
        int m = matrix.length, n = matrix[0].length;
        int lo = Integer.MAX_VALUE, hi = Integer.MIN_VALUE;
        for (int[] row : matrix) {
            lo = Math.min(lo, row[0]);          // smallest possible answer
            hi = Math.max(hi, row[n - 1]);      // biggest possible answer
        }
        int need = (m * n + 1) / 2;             // the median is the need-th smallest

        while (lo < hi) {
            int mid = lo + (hi - lo) / 2;
            int count = 0;
            for (int[] row : matrix) count += countAtMost(row, mid);
            if (count < need) lo = mid + 1;     // too few hands up: the median is bigger
            else hi = mid;                      // enough hands up: the median is mid or smaller
        }
        return lo;
    }

    // how many numbers in a sorted row are <= x (the first index holding a number > x)
    private int countAtMost(int[] row, int x) {
        int left = 0, right = row.length;
        while (left < right) {
            int mid = left + (right - left) / 2;
            if (row[mid] <= x) left = mid + 1;
            else right = mid;
        }
        return left;
    }
}

How long it takes: the guesses halve the range from min to max, so there are about log(max - min) guesses, at most 32 for int. Each guess counts in m rows with an O(log n) search. Total O(m * log(n) * log(max - min)). Extra room: O(1).

Why is the answer always a real number from the grid? Because lo stops at the smallest value whose count reaches need. Any value that is not in the grid has the same count as the grid number just below it, so it cannot be the smallest such value. In the animation, 4 is not in the grid, and the search moves past it.

Watch out for

• (lo + hi) / 2 can overflow when the numbers are huge or very negative. Use lo + (hi - lo) / 2, or long.
• Negative numbers: the range must start at the real minimum, not at 0 or 1.
• lo and hi must come from the first and last columns. The top left and bottom right cells are not the min and max here, because columns are not sorted.
• Off by one in need. For 9 numbers the median is the 5th, so need = (9 + 1) / 2 = 5.

What interviewers ask next

• What if m * n is even? Find the (m*n)/2th and the (m*n)/2 + 1th smallest with the same search (change only need) and average them.
• The kth smallest number in this grid? Same code with need = k.
• Can you use a heap instead? Put the first number of each row in a min heap, pop need times, and after each pop push the next number from the same row. That is O(need * log m) time and O(m) room, which is slower when the grid is big.
• What if the grid is also sorted by columns? Then you can count in O(m + n) with a staircase walk from the bottom left, instead of a halving search in every row.
• Why not binary search on positions? Because positions in the grid say nothing about rank when only rows are sorted. Values do: count of numbers <= x only grows as x grows, and that is all binary search needs.

Remember this

Grown-up words: raising hands for "at most this height" is counting with an upper bound search, guessing and halving the range of values is binary search on the answer, and the middle seat is the median, the (m*n+1)/2th order statistic.