Kth Smallest Element in a Sorted Matrix

Problem Statement

Given an n x n matrix where each row and column is sorted ascending, return the kth smallest element (in overall sorted order).

Example

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

Approach (Binary Search on Value)

Binary search on the value range [matrix[0][0], matrix[n-1][n-1]]. For a candidate mid, count how many elements are <= mid using a staircase walk from the bottom-left in O(n). If the count is < k, raise the low bound; otherwise lower the high bound. The low bound converges to the answer.

Complexity

• Time: O(n log(maxVal - minVal))
• Space: O(1)

Solution Code

Python

def kthSmallest(matrix, k):
    n = len(matrix)
    lo, hi = matrix[0][0], matrix[n-1][n-1]
    while lo < hi:
        mid = (lo + hi) // 2
        count, j = 0, n - 1
        for i in range(n):
            while j >= 0 and matrix[i][j] > mid:
                j -= 1
            count += j + 1
        if count < k:
            lo = mid + 1
        else:
            hi = mid
    return lo

Java

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;
        int count = 0, j = n - 1;
        for (int i = 0; i < n; i++) {
            while (j >= 0 && matrix[i][j] > mid) j--;
            count += j + 1;
        }
        if (count < k) lo = mid + 1;
        else hi = mid;
    }
    return lo;
}

C++

int kthSmallest(vector<vector<int>>& matrix, int k) {
    int n = matrix.size();
    int lo = matrix[0][0], hi = matrix[n-1][n-1];
    while (lo < hi) {
        int mid = lo + (hi - lo) / 2;
        int count = 0, j = n - 1;
        for (int i = 0; i < n; i++) {
            while (j >= 0 && matrix[i][j] > mid) j--;
            count += j + 1;
        }
        if (count < k) lo = mid + 1;
        else hi = mid;
    }
    return lo;
}

Key Takeaways

• Binary search the value, not an index, when elements are not in one sorted line.
• The staircase count of elements <= mid runs in O(n) per step.
• Reference: LeetCode 378 (Kth Smallest Element in a Sorted Matrix).