Problem 18.7 · Heaps and Priority QueuesMedium

Kth Smallest Element in a Sorted Matrix

PatternK-Way Merge Binary Search on the Answer

What it teaches: Rows as sorted lists: K-way merge, or binary search on the value with a staircase count.

Practise it on judges as “Kth Smallest Element in a Sorted Matrix”.

The problem

In an n × n matrix whose rows and columns are sorted ascending, return the k-th smallest value (counting duplicates).

Example 1

Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
Output: 13

Constraints

1 ≤ n ≤ 300
1 ≤ k ≤ n²

Pattern clues in the wording

→ Several sorted sequences
→ k-th in combined order

These clues point to K-Way Merge: Put the first element of each sorted list in a min-heap, repeatedly take the smallest, and push its successor.

Stuck? Take one hint at a time

Solution.java · starter

import java.util.*;

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        return 0;
    }
}

Write your solution locally or in your editor for now. The in-browser runner (Java first, then Python, C++ and more) will run these tests right here.

Test cases

#	Input	Expected
1	matrix = [[1,5,9],[10,11,13],[12,13,15]] k = 8	13
2	matrix = [[-5]] k = 1	-5

+ 1 hidden test the code runner will check

From slow to fast

Approaches

1. K-way merge over rows · O(k log n)2. Binary search on the value · O(n log(max − min))

K-way merge over rows

Time O(k log n) Space O(n)

Push the first element of each row as (value, row, col). Pop k − 1 times, pushing the next element in that row each time. The top is the answer.

Approach 1

import java.util.PriorityQueue;

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> Integer.compare(a[0], b[0]));
        for (int r = 0; r < Math.min(n, k); r++) pq.offer(new int[]{matrix[r][0], r, 0});
        for (int i = 0; i < k - 1; i++) {
            int[] t = pq.poll();
            if (t[2] + 1 < n) pq.offer(new int[]{matrix[t[1]][t[2] + 1], t[1], t[2] + 1});
        }
        return pq.peek()[0];
    }
}

Verdict: Direct use of the pattern.

Binary search on the value

Time O(n log(max − min)) Space O(1)

Search v between matrix[0][0] and matrix[n−1][n−1]. Count entries ≤ v by walking from the bottom-left: if the value is ≤ v, the whole column above counts; move right; else move up. Find the smallest v with count ≥ k.

Approach 2

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length, lo = matrix[0][0], hi = matrix[n - 1][n - 1];
        while (lo < hi) {
            int mid = lo + (hi - lo) / 2;
            if (countAtMost(matrix, mid) >= k) hi = mid; else lo = mid + 1;
        }
        return lo;
    }

    private int countAtMost(int[][] m, int v) {
        int n = m.length, r = n - 1, c = 0, count = 0;
        while (r >= 0 && c < n) {
            if (m[r][c] <= v) { count += r + 1; c++; }
            else r--;
        }
        return count;
    }
}

Verdict: Better for large k.

Before you submit

Edge cases and common mistakes

Test these inputs

k = 1
k = n²
Duplicates

Mistakes people make

In binary search, returning mid when the count equals k (mid might not be in the matrix).

Interview

Follow-up questions

Why does the binary search return a value that's actually in the matrix?

Command Palette