Median in a Row-Wise Sorted Matrix

Given an r × c matrix (an int[][]) where every row is sorted and r · c is odd, return the median of all its values: the element at position (r · c + 1) / 2 in sorted order.

Example

Text

1 3 5
2 6 9
3 6 9

Sorted: 1 2 3 3 5 6 6 9 9  ->  Output: 5 (the 5th of 9)

Brute force: flatten and sort

Java

import java.util.Arrays;

public class Main {
    static int matrixMedian(int[][] matrix) {
        int rows = matrix.length;
        int cols = matrix[0].length;
        int[] all = new int[rows * cols];
        int k = 0;
        for (int[] row : matrix) {
            for (int x : row) {
                all[k++] = x;
            }
        }
        Arrays.sort(all);
        return all[all.length / 2];
    }

    public static void main(String[] args) {
        int[][] m = {{1, 3, 5}, {2, 6, 9}, {3, 6, 9}};
        System.out.println(matrixMedian(m));
    }
}

Output

Time: O(rc log(rc)). Space: O(rc).

Optimal: binary search on the value

The median is the smallest value x such that more than half of the elements are <= x. Counting the elements <= x is cheap because each row is sorted: a binary search per row (countAtMost below, which finds the first index whose value is greater than x) gives the count in O(r log c).

So binary search over values instead of positions. The answer lies between the smallest first element and the largest last element. For each candidate mid, count the elements <= mid. If there are at least half = (rc + 1) / 2, the median is mid or smaller; otherwise it is larger.

Java

public class Main {
    // Number of elements in the sorted row that are <= x.
    static int countAtMost(int[] row, int x) {
        int lo = 0;
        int hi = row.length;
        while (lo < hi) {
            int mid = lo + (hi - lo) / 2;
            if (row[mid] <= x) {
                lo = mid + 1;
            } else {
                hi = mid;
            }
        }
        return lo;
    }

    static int matrixMedian(int[][] matrix) {
        int lo = Integer.MAX_VALUE;
        int hi = Integer.MIN_VALUE;
        for (int[] row : matrix) {
            lo = Math.min(lo, row[0]);
            hi = Math.max(hi, row[row.length - 1]);
        }
        int half = (matrix.length * matrix[0].length + 1) / 2;
        while (lo < hi) {
            int mid = lo + (hi - lo) / 2;
            int atMostMid = 0;
            for (int[] row : matrix) {
                atMostMid += countAtMost(row, mid);
            }
            if (atMostMid >= half) {
                hi = mid;
            } else {
                lo = mid + 1;
            }
        }
        return lo;
    }

    public static void main(String[] args) {
        int[][] m = {{1, 3, 5}, {2, 6, 9}, {3, 6, 9}};
        System.out.println(matrixMedian(m));
        System.out.println(matrixMedian(new int[][] {{1, 3, 4}, {2, 5, 6}, {7, 8, 9}}));
        System.out.println(matrixMedian(new int[][] {{-5, 0, 7}}));
    }
}

Output

5
5
0

Time: O(r log c · log(max - min)): about 32 iterations for int values. Space: O(1).

The search always ends on a value that is in the matrix. lo is the smallest value with enough elements <= lo; if it were not in the matrix, lo - 1 would have the same count and would have been found instead.

hi - lo is computed in int and wraps around when lo is very negative and hi very positive. For inputs that use the full int range, declare lo, hi, and mid as long.

Edge cases

Single row: the middle element of that row.
Duplicates: the count includes all copies, so repeated values are handled.
Even r · c: the problem rules it out. For a lower median, the same search works with half = rc / 2.