☕ DSA Java Articles

Topics › Binary Search › BS on 2D Arrays

Matrix Median

Hard TUF ID 67 A2ZSDE

Java source ↗ Edit this article ↗

Problem

You are given an n × m integer matrix in which every row is sorted in non-decreasing order (columns carry no guarantee), and n × m is odd. Return the median of all n × m values: the value that would sit in the middle if every cell were poured into one list and sorted.

Constraints: 1 ≤ n, m ≤ 10^3, n × m is odd, values fit in a 32-bit signed integer.

Example 1

mat = [[1, 3, 5],
       [2, 6, 9],
       [3, 6, 9]]
answer = 5           // sorted: 1 2 3 3 5 6 6 9 9, the middle (5th) value is 5

Example 2

mat = [[1, 3, 8],
       [2, 3, 4],
       [1, 2, 5]]
answer = 3           // sorted: 1 1 2 2 3 3 4 5 8

Example 3

mat = [[1, 2, 3, 4, 5]]
answer = 3           // a single sorted row: its middle element

With total = n × m cells, the median is the element at index total / 2 of the sorted list (0-indexed). A useful way to restate that: the median is the smallest value x such that more than total / 2 cells are <= x. Both binary search approaches below are built on that restatement.

Approach 1: Brute force, flatten and sort

Intuition. Copy all cells into one array, sort it, read the middle. This works for any matrix and does not use the sorted rows at all.

Algorithm.

  1. Copy every cell into an array of length n × m.
  2. Sort it.
  3. Return the element at index (n × m) / 2.
static int bruteForce(int[][] mat) {
    int n = mat.length, m = mat[0].length;
    int[] all = new int[n * m];
    int k = 0;
    for (int[] row : mat) for (int x : row) all[k++] = x;
    Arrays.sort(all);
    return all[(n * m) / 2];
}

Complexity. Time O(n m log(n m)) for the sort. Space O(n m) for the copy.

For 10^6 cells this is a few tens of millions of operations and a megabyte-scale copy: workable, but the sorted rows let us avoid both the sort and the copy.

Approach 2: Better, binary search on the value with a linear count

Intuition. Turn the question around. Instead of “which element is the median?”, ask “given a value x, how many cells are <= x?” Call that count(x). As x grows, count(x) can only grow, so the predicate “count(x) > total / 2” is false for small x and true from the median onward. The median is the smallest x for which it is true, and binary search over the value range finds it.

The range to search is [smallest cell, largest cell]. Since rows are sorted, the smallest cell is the minimum of the first column and the largest is the maximum of the last column. In this approach count(x) is computed naively by scanning every cell; the point is to see the search structure before optimising the count.

Algorithm.

  1. Set lo to the minimum first element over all rows and hi to the maximum last element.
  2. Let half = total / 2.
  3. While lo <= hi: compute mid and count(mid).
    • If count(mid) <= half, too few cells are <= mid, so the median is larger: lo = mid + 1.
    • Otherwise mid is at least the median: hi = mid − 1.
  4. Return lo.
static int countLinear(int[][] mat, int x) {
    int count = 0;
    for (int[] row : mat) for (int v : row) if (v <= x) count++;
    return count;
}

static int better(int[][] mat) {
    int n = mat.length, m = mat[0].length;
    int lo = Integer.MAX_VALUE, hi = Integer.MIN_VALUE;
    for (int[] row : mat) { lo = Math.min(lo, row[0]); hi = Math.max(hi, row[m - 1]); }
    int half = (n * m) / 2;
    while (lo <= hi) {
        int mid = (int) Math.floorDiv((long) lo + hi, 2);   // long: lo + hi overflows int when the range is wide
        if (countLinear(mat, mid) <= half) lo = mid + 1;    // too few values <= mid: the median is larger
        else hi = mid - 1;                                  // mid already covers the middle position
    }
    return lo;
}

Complexity. Time O(n m log(range)), where range is the spread of values, at most 2^32, so about 32 full scans. Space O(1).

Why does the loop end at the median? Let M be the median. For any x < M, at most half cells are <= x (the cells at sorted positions 0 .. half − 1), so the test fails and lo moves above x. For x >= M, at least half + 1 cells are <= x, so the test passes and hi moves below x. The two pointers close in on M from both sides and lo finishes at exactly M, which is a real cell value because the search range contains it.

Approach 3: Optimal, binary search on the value with per-row upper bounds

Intuition. The only expensive part of Approach 2 is count(x). Each row is sorted, so “how many values in this row are <= x” is the upper bound of x: the index of the first element greater than x, found by binary search in O(log m). Summing the upper bounds over all rows gives the full count in O(n log m) instead of O(n m).

Algorithm.

  1. Same search range and loop as Approach 2.
  2. For each probe mid, compute count = sum over rows of upperBound(row, mid).
  3. Move lo or hi exactly as before; return lo.
static int upperBound(int[] row, int x) {
    int lo = 0, hi = row.length - 1, ans = row.length;
    while (lo <= hi) {
        int mid = lo + (hi - lo) / 2;
        if (row[mid] > x) { ans = mid; hi = mid - 1; }
        else lo = mid + 1;
    }
    return ans;
}

static int optimal(int[][] mat) {
    int n = mat.length, m = mat[0].length;
    int lo = Integer.MAX_VALUE, hi = Integer.MIN_VALUE;
    for (int[] row : mat) { lo = Math.min(lo, row[0]); hi = Math.max(hi, row[m - 1]); }
    int half = (n * m) / 2;
    while (lo <= hi) {
        int mid = (int) Math.floorDiv((long) lo + hi, 2);
        int count = 0;
        for (int[] row : mat) count += upperBound(row, mid);
        if (count <= half) lo = mid + 1;
        else hi = mid - 1;
    }
    return lo;
}

Dry run on Example 1. total = 9, half = 4, lo = min(1, 2, 3) = 1, hi = max(5, 9, 9) = 9.

lo hi mid upper bounds per row count count > 4? action
1 9 5 3, 1, 1 5 yes hi = 4
1 4 2 1, 1, 0 2 no lo = 3
3 4 3 2, 1, 1 4 no lo = 4
4 4 4 2, 1, 1 4 no lo = 5

lo = 5 > hi = 4, so the median is 5. Notice how mid = 4 is not a cell value, yet the search handles it fine: the count simply does not change between 3 and 4, and lo keeps moving up until it reaches the smallest value whose count exceeds half.

Complexity. Time O(n log m log(range)): about 32 probes, each doing n binary searches of log m steps. For n = m = 1000 that is roughly 32 × 1000 × 10 = 320,000 operations, against 10^6 × 20 for sorting. Space O(1).

Edge cases and pitfalls