LeetCode 378 - Kth Smallest Element in a Sorted Matrix

The problem gives us an n x n matrix where both rows and columns are sorted in ascending order. This means two ordering guarantees exist simultaneously: - Every row is sorted from left to right. - Every column is sorted from top to bottom.

leetcodemediumarraybinary-searchsortingheap-(priority-queue)matrix

LeetCode Problem 378

Difficulty: 🟡 Medium
Topics: Array, Binary Search, Sorting, Heap (Priority Queue), Matrix

Solution

Problem Understanding

The problem gives us an n x n matrix where both rows and columns are sorted in ascending order. This means two ordering guarantees exist simultaneously:

  • Every row is sorted from left to right.
  • Every column is sorted from top to bottom.

We are asked to return the kth smallest element across the entire matrix when all elements are considered together in sorted order.

The important detail is that we are looking for the kth smallest element overall, not the kth distinct value. Duplicate values count separately. For example, if the matrix contains two 13s, both participate independently in the ordering.

For the first example:

[
  [1, 5, 9],
  [10,11,13],
  [12,13,15]
]

If we flatten and sort all values, we get:

[1,5,9,10,11,12,13,13,15]

The 8th smallest element is 13.

The constraints are important because they strongly influence which algorithms are feasible:

  • n can be as large as 300
  • The matrix can therefore contain up to 90,000 elements
  • Values may be negative and may range up to 10^9
  • Memory usage must be better than O(n^2)

A naive solution that copies every value into another array and sorts it works, but it ignores the matrix's sorted structure. The problem is specifically designed so we can exploit that structure for a more efficient solution.

Several edge cases are important:

  • A 1 x 1 matrix, where the only element must be returned.
  • Duplicate values, since the problem counts duplicates separately.
  • Negative numbers, which can expose incorrect binary search boundaries.
  • Cases where k = 1, meaning we need the global minimum.
  • Cases where k = n^2, meaning we need the global maximum.

The matrix is guaranteed to be square and sorted correctly, so we do not need to validate the input structure.

Approaches

Brute Force Approach

The simplest solution is to flatten the matrix into a single array, sort the array, and return the element at index k - 1.

The algorithm is straightforward:

  1. Traverse every cell in the matrix.
  2. Store all values in a list.
  3. Sort the list.
  4. Return the k - 1 indexed element.

This works because sorting produces the exact global ordering of all matrix elements.

However, this approach ignores the fact that the matrix is already partially sorted. Since the matrix can contain up to 90,000 elements, sorting all elements costs O(n^2 log(n^2)), which is unnecessarily expensive.

It also uses O(n^2) extra memory, violating the follow-up requirement for better memory usage.

Optimal Approach, Binary Search on Value Range

The key insight is that the matrix is sorted both row-wise and column-wise. This allows us to efficiently determine how many numbers are less than or equal to a chosen value.

Instead of binary searching indices, we binary search the numeric value range itself.

The smallest possible answer is:

matrix[0][0]

The largest possible answer is:

matrix[n-1][n-1]

Suppose we pick some value mid.

If we can efficiently count how many elements in the matrix are <= mid, then:

  • If fewer than k elements are <= mid, the answer must be larger.
  • Otherwise, the answer could be mid or smaller.

This creates a classic monotonic condition suitable for binary search.

The crucial optimization is counting elements efficiently in O(n) time instead of scanning all n^2 cells.

Because rows and columns are sorted, we can start from the bottom-left corner:

  • If the current value is <= mid, then everything above it in that column is also <= mid.
  • If the current value is > mid, move upward.

This counting process runs in linear time per binary search iteration.

Approach Comparison

Approach Time Complexity Space Complexity Notes
Brute Force O(n^2 log(n^2)) O(n^2) Flatten and sort all elements
Optimal O(n log(max-min)) O(1) Binary search over value range with linear counting

Algorithm Walkthrough

Optimal Algorithm, Binary Search on Values

  1. Initialize the binary search boundaries using the smallest and largest values in the matrix.

The smallest element is always at matrix[0][0] because rows and columns are sorted ascending.

The largest element is always at matrix[n-1][n-1]. 2. Perform binary search over the numeric range.

Compute:

mid = (left + right) // 2

We now need to determine how many elements in the matrix are less than or equal to mid. 3. Count elements <= mid using the bottom-left traversal technique.

Start from:

row = n - 1
col = 0

At each step:

  • If matrix[row][col] <= mid, then every element above it in the same column is also <= mid.

Therefore, add:

row + 1

to the count, then move right.

  • Otherwise, move upward because the current value is too large.
  1. Compare the count against k.
  • If count < k, then the answer must be larger than mid.

Move:

left = mid + 1
  • Otherwise, the answer is mid or smaller.

Move:

right = mid
  1. Continue until left == right.

At that point, the binary search has converged to the smallest value whose rank in the matrix is at least k. 6. Return left.

Why it works

The algorithm relies on a monotonic property.

For any value x, define:

count(x) = number of elements <= x

As x increases, count(x) never decreases.

This means we can binary search for the smallest value such that:

count(x) >= k

That value is exactly the kth smallest element.

The matrix ordering guarantees that counting can be done efficiently in linear time.

Python Solution

from typing import List

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)

        def count_less_equal(target: int) -> int:
            row = n - 1
            col = 0
            count = 0

            while row >= 0 and col < n:
                if matrix[row][col] <= target:
                    count += row + 1
                    col += 1
                else:
                    row -= 1

            return count

        left = matrix[0][0]
        right = matrix[n - 1][n - 1]

        while left < right:
            mid = (left + right) // 2

            if count_less_equal(mid) < k:
                left = mid + 1
            else:
                right = mid

        return left

The implementation begins by defining a helper function named count_less_equal. This function determines how many matrix elements are less than or equal to a target value.

Instead of scanning every element, it starts from the bottom-left corner. This position is useful because moving upward decreases values while moving right increases values.

When the current value is less than or equal to the target, the entire column section above that value is also valid because columns are sorted. We can therefore add row + 1 elements at once.

The main algorithm performs binary search over the value range rather than over indices.

left starts at the global minimum value and right starts at the global maximum value.

For each midpoint:

  • If fewer than k elements are less than or equal to mid, the answer must be larger.
  • Otherwise, the answer lies at mid or below.

Eventually, the search space collapses to a single value, which is the kth smallest element.

Go Solution

func kthSmallest(matrix [][]int, k int) int {
	n := len(matrix)

	countLessEqual := func(target int) int {
		row := n - 1
		col := 0
		count := 0

		for row >= 0 && col < n {
			if matrix[row][col] <= target {
				count += row + 1
				col++
			} else {
				row--
			}
		}

		return count
	}

	left := matrix[0][0]
	right := matrix[n-1][n-1]

	for left < right {
		mid := left + (right-left)/2

		if countLessEqual(mid) < k {
			left = mid + 1
		} else {
			right = mid
		}
	}

	return left
}

The Go implementation follows the same algorithmic structure as the Python version.

One small difference is the midpoint calculation:

mid := left + (right-left)/2

This avoids potential integer overflow, which is a common defensive practice in Go and other statically typed languages.

Go slices are used directly for matrix access, and no additional memory allocations are required beyond a few integer variables, preserving the O(1) extra space complexity.

Worked Examples

Example 1

Input:

matrix = [
  [1, 5, 9],
  [10,11,13],
  [12,13,15]
]

k = 8

Initial boundaries:

Variable Value
left 1
right 15

Iteration 1

mid = (1 + 15) // 2 = 8

Count elements <= 8.

Traversal:

Position Value Action Count
(2,0) 12 move up 0
(1,0) 10 move up 0
(0,0) 1 add 1, move right 1
(0,1) 5 add 1, move right 2
(0,2) 9 move up 2

Final count:

2

Since 2 < 8, move right:

left = 9

Iteration 2

mid = (9 + 15) // 2 = 12

Count elements <= 12.

Values are:

1,5,9,10,11,12

Count:

6

Since 6 < 8:

left = 13

Iteration 3

mid = (13 + 15) // 2 = 14

Elements <= 14:

1,5,9,10,11,12,13,13

Count:

8

Since 8 >= 8:

right = 14

Iteration 4

mid = (13 + 14) // 2 = 13

Count elements <= 13:

8

Since 8 >= 8:

right = 13

Now:

left == right == 13

Return:

13

Example 2

Input:

matrix = [[-5]]
k = 1

Initial values:

Variable Value
left -5
right -5

The binary search loop never runs because:

left == right

Return:

-5

Complexity Analysis

Measure Complexity Explanation
Time O(n log(max-min)) Binary search iterations multiplied by O(n) counting
Space O(1) Only a few variables are used

The counting function runs in O(n) time because each move either decreases the row index or increases the column index. No cell is visited more than once during a count operation.

The binary search range spans from the smallest to the largest matrix value. If the numeric range is max - min, the binary search requires:

O(log(max - min))

iterations.

Multiplying both parts gives:

O(n log(max - min))

overall time complexity.

The algorithm uses constant extra memory because it does not allocate auxiliary arrays, heaps, or hash maps.

Test Cases

sol = Solution()

# Example case from the prompt
assert sol.kthSmallest(
    [[1,5,9],[10,11,13],[12,13,15]],
    8
) == 13  # standard mixed matrix

# Single element matrix
assert sol.kthSmallest([[-5]], 1) == -5  # smallest possible matrix

# Matrix with duplicates
assert sol.kthSmallest(
    [[1,2],[2,3]],
    2
) == 2  # duplicate values counted separately

# k equals 1
assert sol.kthSmallest(
    [[1,3],[5,7]],
    1
) == 1  # global minimum

# k equals n^2
assert sol.kthSmallest(
    [[1,3],[5,7]],
    4
) == 7  # global maximum

# Negative numbers
assert sol.kthSmallest(
    [[-10,-5],[-4,0]],
    3
) == -4  # negative values handled correctly

# Larger duplicate-heavy matrix
assert sol.kthSmallest(
    [[1,1,3],[1,2,2],[2,3,3]],
    5
) == 2  # duplicate-heavy ordering

# All elements identical
assert sol.kthSmallest(
    [[7,7],[7,7]],
    3
) == 7  # repeated identical values

# Increasing rows and columns
assert sol.kthSmallest(
    [[1,2,3],[4,5,6],[7,8,9]],
    6
) == 6  # perfectly ordered matrix

Test Case Summary

Test Why
Standard example Verifies normal functionality
Single element matrix Validates smallest input size
Duplicate values Ensures duplicates are counted correctly
k = 1 Confirms minimum element retrieval
k = n^2 Confirms maximum element retrieval
Negative values Verifies binary search boundaries with negatives
Duplicate-heavy matrix Tests repeated ordering behavior
All identical values Ensures binary search handles equal ranges
Fully increasing matrix Validates standard sorted structure

Edge Cases

Single Element Matrix

A 1 x 1 matrix is the smallest valid input. A careless implementation might incorrectly assume multiple rows or columns exist.

For example:

[[-5]]

The implementation handles this naturally because the binary search boundaries both become -5, so the loop terminates immediately and returns the only value.

Duplicate Values

The problem explicitly states that duplicates count separately. This is important because some incorrect solutions accidentally treat values as unique.

For example:

[
  [1,2],
  [2,3]
]

The sorted ordering is:

[1,2,2,3]

The 2nd and 3rd smallest elements are both 2.

The counting function correctly counts every occurrence independently, so duplicates are handled automatically.

Negative Numbers

Negative values can expose bugs in binary search implementations, especially if boundaries are initialized incorrectly.

For example:

[
  [-10,-5],
  [-4,0]
]

The algorithm safely initializes:

left = matrix[0][0]
right = matrix[n-1][n-1]

Since these are actual matrix values, the search space always remains valid regardless of whether numbers are positive or negative.

All Elements Equal

When every value is identical, many binary search implementations risk infinite loops if boundary updates are incorrect.

Example:

[
  [7,7],
  [7,7]
]

The implementation avoids this issue because:

  • mid always equals 7
  • right = mid eventually collapses the search range
  • The loop condition uses left < right

This guarantees termination.