378. Kth Smallest Element in a Sorted Matrix

At first glance, this problem looks like a simple sorting task, but the real learning begins when you try to do it without extra memory. It makes you think about how to use binary search in a slightly different way. In fact, this is a known favorite at Google, Amazon, Meta, Microsoft, and Bloomberg, so it’s worth practicing until the idea feels natural.

Problem statement

You are given an $n \times n$ matrix where each row and each column is sorted in ascending order. Return the $k^{th}$ smallest element in the matrix.

Constraints:

• n $==$ matrix.length $==$ matrix[i].length
• $1 \leq$ n $\leq 300$
• $−10^9 \leq$ matrix[i][j] $\leq 10^9$
• All the rows and columns of matrix are guaranteed to be sorted in non-decreasing order.
• $1 \leq$ k $\leq n^2$

Examples

A Naive start

A first thought might be to solve the problem by flattening the entire matrix into a single list, sorting it, and then picking the $k^{th}$ element. This requires storing all $n^2$ elements in memory, which takes $O(n^2)$ space, and then sorting them takes $O(n^2 \log n)$ time. However, as the problem explicitly asks for a solution with memory usage better than $O(n^2)$, we need something better.

An optimized way: Using the Binary Search

After realizing we can’t afford to store and sort all $n^2$ values, the next question becomes: Can we find the $k^{th}$ smallest without fully ordering every element of the matrix?

The key observation is that the matrix is already partially ordered because every row and column is sorted. This means we don’t need to guess the answer by position, we can guess it by value. We know that the smallest value is at the top-left cell matrix[0][0] and the largest value is at the bottom-right cell matrix[n−1], so the $k^{th}$ smallest element must lie somewhere in the numeric range [matrix[0][0],matrix[n−1][n−1]].

Now, this allows us to find the $k^{th}$ smallest element efficiently using Binary Search on the value range. Instead of searching by index, we search by value. For any guessed value $x$, we count how many elements in the matrix are $\leq x$. If this count is less than $k$, $x$ is too small, so we search higher values. Otherwise, $x$ is big enough and we keep searching smaller values. Repeating this shrinks the value range until it collapses to a single number, which is the $k^{th}$ smallest element.

While binary search is efficient, we also need an efficient way to count how many elements are $\leq x$ without scanning all $n^2$ cells. For this, we use a staircase walk starting from the bottom-left corner of the matrix. If the current value is $\le x$, then everything above it in the same column is also $\le x$, so we add all those elements at once and move right. If the current value is $>x$, we move up to reach smaller values. As we only move up or right, we visit at most about $2n$ cells, not $n^2$.

Step-by-step algorithm

Let’s look at the algorithm steps of the solution:

1. Set up the search range (by value): n = len(matrix), low = matrix[0][0], and high = matrix[n - 1][n - 1].
2. Start the binary search on the value range. While low < high:
   1. Compute mid = (low + high) // 2.
   2. Compute the number of elements that are <= mid, i.e., num_count = countLeq(matrix, mid).
      1. Inside countLeq, initialize row = n - 1 (bottom row), col = 0 (leftmost column), and num_count = 0.
      2. While row >= 0 and col < n:
         1. Update current = matrix[row][col].
         2. If current <= x, add all values above in this column, i.e., num_count += row + 1, and move to the right col += 1.
         3. Otherwise, move upward row -= 1.
      3. Return num_count once the loop ends.
   3. Shrink the search range based on num_count.
      1. If num_count < k, the k-th smallest is larger than mid, so set low = mid + 1.
      2. Otherwise, the k-th smallest is mid or smaller, so set high = mid.
3. When low == high, the loop ends, so return low as the final answer.

Here’s an illustration to help you visualize the solution:

1 / 6

Code implementation in Python

Now, let’s look at the code for this solution below:

from typing import List

def kthSmallest(matrix: List[List[int]], k: int) -> int:
    n = len(matrix)
    low = matrix[0][0]
    high = matrix[n - 1][n - 1]

    # We binary search in the VALUE range [low, high]
    while low < high:

        # Pick the middle value in the current search range
        mid = (low + high) // 2

        # Count how many values in the matrix are <= mid
        num_count = countLeq(matrix, mid)

        # If there are fewer than k values <= mid,
        # then the k-th smallest must be larger than mid.
        if num_count < k:
            # So we move the lower bound up (excluding mid)
            low = mid + 1

        else:
            # Otherwise, mid is big enough (could be the answer or too large),
            # so we keep mid and shrink the upper bound down to mid.
            high = mid

    # When low == high, we've found the smallest value that has at least k numbers <= it
    return low
    

def countLeq(matrix: List[List[int]], x: int) -> int:
    n = len(matrix)

    # Start from bottom-left corner.
    row = n - 1
    col = 0

    # This will store how many elements are <= x.
    num_count = 0

    # Keep moving until we walk out of the matrix boundaries.
    while row >= 0 and col < n:

        # Current cell value we are standing on
        current = matrix[row][col]

        # If current value is <= x, then ALL values above it in the same column.
        if current <= x:

            # Add that whole column chunk in one go.
            num_count += row + 1

            # Move right to the next column (which may have larger values).
            col += 1

        else:
            # If current value is > x, move up to find smaller values.
            row -= 1

    # Return the total number of elements <= x
    return num_count


# Driver code
def main():
    # Each test case is (matrix, k)
    test_cases = [
        ([[7]], 1),

        ([[1, 2],
          [2, 4]], 3),

        ([[-8, -3,  1],
          [-6, -1,  2],
          [-6,  0,  5]], 4),

        ([[1,  3,  3,  6],
          [2,  4,  4,  7],
          [3,  5,  5,  8],
          [6,  8,  9, 10]], 9),

        ([[-10, -4, -4,  0,  3],
          [ -8, -3,  1,  2,  6],
          [ -7, -2,  2,  4,  8],
          [ -5,  0,  3,  7,  9],
          [ -1,  2,  6,  8, 12]], 17),
    ]

    # Run each test case one by one
    for i, (matrix, k) in enumerate(test_cases, 1):

        # Compute k-th smallest for this matrix
        result = kthSmallest(matrix, k)

        # Print the input and output clearly
        print(f"{i}.\tmatrix = {matrix}")
        print(f"\tk = {k}")
        print(f"\n\tOutput: {result}")
        print("-" * 75)


if __name__ == "__main__":
    main()

Time complexity

• The binary search step takes $O(log⁡(maxValue−minValue)$.
  • We binary search over the value range from the smallest to the largest element. Each iteration cuts the range in half, so it takes about $O(\log_{2} (maxValue – minValue))$ iterations.
• The counting step takes $O(n)$.
  • We start at the bottom-left and move only up or right. In the worst case, we can move up at most $n$ times and right at most $n$ times, at most about $2n$ moves in total. So, one ‘count how many ≤ x’ pass is linear in $n$.

The overall time complexity is $O(n \times \log(maxValue − minValue))$.

Space complexity

The space complexity of this solution is $O(1)$, meaning it uses constant extra memory.

Common mistakes to avoid

This problem is straightforward once it clicks, but small implementation mistakes are very common. Keep an eye for these common mistakes:

• Binary searching indices instead of values: The matrix is sorted, but not fully sorted as a flat array. Binary search works here because we search the value range, not positions.
• Flattening the matrix and sorting: This violates the space constraint (O(n²) space) and is often an instant rejection in interviews.
• Incorrect counting logic: A very common bug is:
  • Forgetting to add row + 1 when matrix[row][col] <= mid.
  • Or, starting the staircase walk from the wrong corner (top-left does NOT work; bottom-left or top-right does).
• Off-by-one errors in range updates
  • Using high = mid - 1 instead of high = mid.
  • Or, returning mid instead of the final low.