Kth smallest element in a row-column wise sorted matrix

The problem, simply

You get an n by n grid. Every row grows from left to right, and every column grows from top to bottom. You also get a number k.

Give back the kth smallest number in the whole grid. Repeated numbers count separately, so in 1, 2, 2, 3 the 3rd smallest is 2.

A story to remember it

Picture it

Examples

Way 1: the simple way

Pour all numbers into one line, sort it, and take seat k - 1.

import java.util.*;

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        int[] all = new int[n * n];
        int i = 0;
        for (int[] row : matrix)
            for (int value : row) all[i++] = value;
        Arrays.sort(all);
        return all[k - 1];          // the kth smallest sits at index k - 1
    }
}

How long it takes: sorting n * n numbers, O(n^2*log(n)). Extra room: O(n^2) for the copy.

Way 2: the clever way

Binary search on the value, like the teacher's guessing game.

1. lo = matrix[0][0] (the smallest) and hi = matrix[n-1][n-1] (the biggest).
2. While lo < hi, guess mid halfway between them.
3. Count how many numbers are <= mid with a staircase walk from the bottom left corner:
   • If the number is <= mid, everything above it in that column is smaller still. Add row + 1 and step right.
   • Otherwise step up.
4. If the count is less than k, the answer is bigger: lo = mid + 1. Otherwise the answer is mid or smaller: hi = mid.
5. When lo == hi, that is the answer.

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        int lo = matrix[0][0], hi = matrix[n - 1][n - 1];
        while (lo < hi) {
            int mid = lo + (hi - lo) / 2;
            if (countAtMost(matrix, mid) < k) lo = mid + 1;   // too few: answer is bigger
            else hi = mid;                                    // enough: answer is mid or smaller
        }
        return lo;
    }

    // staircase from the bottom left: how many numbers are <= x
    private int countAtMost(int[][] matrix, int x) {
        int n = matrix.length;
        int row = n - 1, col = 0, count = 0;
        while (row >= 0 && col < n) {
            if (matrix[row][col] <= x) {
                count += row + 1;   // this cell and everything above it in the column
                col++;
            } else {
                row--;
            }
        }
        return count;
    }
}

How long it takes: each count walks at most 2n steps, O(n). The guesses halve the range, so there are about log(max - min) of them, at most 32 for int. Total O(n*log(max-min)). Extra room: O(1).

The answer is always a real number in the grid. lo stops at the smallest value with at least k numbers at or below it, and a value that is not in the grid always has the same count as the grid number just below it.

Another way: a heap of row fronts

Treat each row as a sorted line, and merge the lines with a min heap, like the "print in sorted order" problem. Pop the smallest front k times.

import java.util.*;

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> Integer.compare(a[0], b[0]));
        for (int r = 0; r < n; r++) heap.offer(new int[]{matrix[r][0], r, 0});   // {value, row, col}

        for (int i = 1; i < k; i++) {               // throw away the k - 1 smallest
            int[] top = heap.poll();
            int r = top[1], c = top[2];
            if (c + 1 < n) heap.offer(new int[]{matrix[r][c + 1], r, c + 1});
        }
        return heap.poll()[0];                      // the kth pop is the answer
    }
}

This takes O(k*log(n)) time and O(n) room. It is the better pick when k is small. If you only put min(n, k) rows into the heap at the start, it becomes O(k*log(min(n, k))).

Watch out for

• Counting < mid instead of <= mid. The search then answers one more than the right number, 14 instead of 13 here.
• (lo + hi) / 2 overflows for huge values. Use lo + (hi - lo) / 2.
• Starting the staircase at the top left or bottom right. From there both moves change the number the same way.
• In the heap way, (a, b) -> a[0] - b[0] can overflow. Use Integer.compare.

What interviewers ask next

• Which way is better? Binary search uses O(1) room and does not care about k. The heap is simpler and faster when k is tiny compared to n * n.
• Kth largest? It is the (n*n - k + 1)th smallest.
• What if only rows are sorted, not columns? The staircase count no longer works. Count with a halving search in each row instead, O(n log n) per guess. That is the median problem.
• Kth smallest pair sum from two sorted arrays, or kth smallest in a multiplication table. These are hidden sorted grids: the cell (i, j) is a[i] + b[j] or i * j. The same guess and count trick works without building the grid.
• Is there an O(n) method? Yes, the Frederickson and Johnson selection algorithm, but nobody expects you to code it. Mentioning it shows depth.

Remember this

Grown-up words: the grid is a row and column sorted matrix (a Young tableau), the guessing game is binary search on the answer, the walk that counts a column at once is a staircase count, and the pile of row fronts is a min heap doing a k way merge.