Kth Smallest Element in Sorted Matrix

Find kth smallest element in row and column sorted matrix using heap

Language Selection

Choose your preferred programming language

Showing: Python

Kth Smallest Element in Sorted Matrix

Problem Statement

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

Note that it is the kth smallest element in the sorted order, not the kth distinct element.

Input/Output Specifications

  • Input:
    • matrix: n x n matrix with sorted rows and columns
    • k: Integer (1 <= k <= n²)
  • Output: Integer representing kth smallest element

Constraints

  • n == matrix.length == matrix[i].length
  • 1 <= n <= 300
  • -10^9 <= matrix[i][j] <= 10^9
  • All rows and columns are sorted in non-decreasing order
  • 1 <= k <= n²

Examples

Example 1:

Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
Output: 13
Explanation: Elements in sorted order: [1,5,9,10,11,12,13,13,15]
The 8th smallest element is 13.

Example 2:

Input: matrix = [[-5]], k = 1
Output: -5

Solution Approaches

Approach 1: Min-Heap (Optimal for Small k)

Algorithm Explanation:

  1. Start with top-left element in min-heap
  2. For each extraction:
    • Pop minimum element from heap
    • Add right and down neighbors if not visited
    • Use set to avoid duplicates
  3. Return kth extracted element

Implementation:

Python:

import heapq

def kth_smallest_matrix_heap(matrix, k):
    """
    Find kth smallest in sorted matrix using min-heap
    Time: O(min(k, n) * log(min(k, n)))
    Space: O(min(k, n))
    """
    if not matrix or not matrix[0]:
        return -1
    
    n = len(matrix)
    heap = [(matrix[0][0], 0, 0)]  # (value, row, col)
    visited = {(0, 0)}
    
    for _ in range(k):
        val, row, col = heapq.heappop(heap)
        
        if _ == k - 1:  # kth element
            return val
        
        # Add right neighbor
        if col + 1 < n and (row, col + 1) not in visited:
            heapq.heappush(heap, (matrix[row][col + 1], row, col + 1))
            visited.add((row, col + 1))
        
        # Add down neighbor
        if row + 1 < n and (row + 1, col) not in visited:
            heapq.heappush(heap, (matrix[row + 1][col], row + 1, col))
            visited.add((row + 1, col))
    
    return -1

# Optimized version - add first row to heap initially
def kth_smallest_matrix_heap_v2(matrix, k):
    """
    Optimized heap approach starting with first row
    Time: O(k * log n)
    Space: O(n)
    """
    n = len(matrix)
    
    # Initialize heap with first row
    heap = [(matrix[0][j], 0, j) for j in range(n)]
    heapq.heapify(heap)
    
    for _ in range(k):
        val, row, col = heapq.heappop(heap)
        
        if _ == k - 1:
            return val
        
        # Add element from next row in same column
        if row + 1 < n:
            heapq.heappush(heap, (matrix[row + 1][col], row + 1, col))
    
    return -1

Java:

import java.util.*;

class Solution {
    /**
     * Find kth smallest using min-heap
     * Time: O(min(k, n) * log(min(k, n)))
     * Space: O(min(k, n))
     */
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        
        PriorityQueue<int[]> heap = new PriorityQueue<>(
            (a, b) -> Integer.compare(a[0], b[0])
        );
        
        // Add first element
        heap.offer(new int[]{matrix[0][0], 0, 0});
        Set<String> visited = new HashSet<>();
        visited.add("0,0");
        
        for (int i = 0; i < k; i++) {
            int[] current = heap.poll();
            int val = current[0];
            int row = current[1];
            int col = current[2];
            
            if (i == k - 1) {
                return val;
            }
            
            // Add right neighbor
            if (col + 1 < n && !visited.contains(row + "," + (col + 1))) {
                heap.offer(new int[]{matrix[row][col + 1], row, col + 1});
                visited.add(row + "," + (col + 1));
            }
            
            // Add down neighbor
            if (row + 1 < n && !visited.contains((row + 1) + "," + col)) {
                heap.offer(new int[]{matrix[row + 1][col], row + 1, col});
                visited.add((row + 1) + "," + col);
            }
        }
        
        return -1;
    }
}

Approach 2: Binary Search (Optimal for Large k)

Algorithm Explanation:

  1. Binary search on value range [matrix[0][0], matrix[n-1][n-1]]
  2. For each candidate value, count elements ≤ candidate
  3. Find smallest value with at least k elements ≤ it

Implementation:

Python:

def kth_smallest_matrix_binary_search(matrix, k):
    """
    Find kth smallest using binary search on values
    Time: O(n * log(max - min))
    Space: O(1)
    """
    n = len(matrix)
    left, right = matrix[0][0], matrix[n-1][n-1]
    
    def count_less_equal(target):
        count = 0
        row, col = n - 1, 0  # Start from bottom-left
        
        while row >= 0 and col < n:
            if matrix[row][col] <= target:
                count += row + 1  # All elements in column up to row
                col += 1
            else:
                row -= 1
        
        return count
    
    while left < right:
        mid = (left + right) // 2
        
        if count_less_equal(mid) < k:
            left = mid + 1
        else:
            right = mid
    
    return left

# Alternative counting approach
def count_smaller_equal(matrix, target):
    """
    Count elements <= target in sorted matrix
    Time: O(n)
    """
    n = len(matrix)
    count = 0
    row, col = 0, n - 1  # Start from top-right
    
    while row < n and col >= 0:
        if matrix[row][col] <= target:
            count += col + 1  # All elements in row up to col
            row += 1
        else:
            col -= 1
    
    return count

Key Insights

  • Heap Approach: Efficient for small k, explores only necessary elements
  • Binary Search: Efficient for large k, searches on value space
  • Matrix Properties: Utilize sorted rows and columns for optimization
  • Counting Strategy: Can count elements ≤ value in O(n) time

Edge Cases

  1. k = 1: Return top-left element (smallest)
  2. k = n²: Return bottom-right element (largest)
  3. Single Element: Matrix with one element
  4. Duplicate Values: Handle multiple occurrences correctly
  5. All Same Values: All elements identical

Test Cases

def test_kth_smallest_matrix():
    # Example 1
    matrix1 = [[1,5,9],[10,11,13],[12,13,15]]
    assert kth_smallest_matrix_heap(matrix1, 8) == 13
    assert kth_smallest_matrix_binary_search(matrix1, 8) == 13
    
    # Single element
    assert kth_smallest_matrix_heap([[-5]], 1) == -5
    
    # k = 1 (smallest)
    matrix2 = [[1,2],[1,3]]
    assert kth_smallest_matrix_heap(matrix2, 1) == 1
    
    # k = n² (largest)
    assert kth_smallest_matrix_heap(matrix2, 4) == 3
    
    # All same values
    matrix3 = [[1,1],[1,1]]
    assert kth_smallest_matrix_heap(matrix3, 2) == 1
    
    # Negative values
    matrix4 = [[-5,-4],[-3,-2]]
    assert kth_smallest_matrix_heap(matrix4, 2) == -4
    
    print("All tests passed!")

test_kth_smallest_matrix()

Interview Tips

  1. Ask About Matrix Size: Determines best approach (heap vs binary search)
  2. Clarify k: 1-indexed, kth smallest not kth distinct
  3. Utilize Properties: Explain how sorted rows/columns help
  4. Compare Approaches: Discuss time complexity trade-offs
  5. Counting Technique: Show how to count elements ≤ target in O(n)

Follow-up Questions

  1. Kth Largest: Find kth largest element instead
  2. Rectangle Query: Find kth smallest in submatrix
  3. Multiple Queries: Handle multiple k values efficiently
  4. Streaming Updates: Handle matrix updates dynamically
  5. Sparse Matrix: Optimize for matrices with many zeros

Common Mistakes

  1. Wrong Initialization: Not starting heap with correct elements
  2. Duplicate Handling: Adding same position multiple times to heap
  3. Binary Search Bounds: Incorrect search range or termination
  4. Counting Errors: Wrong implementation of element counting
  5. Edge Cases: Not handling k=1 or k=n² properly

Concept Explanations

Matrix Exploration with Heap: Since rows and columns are sorted, we can explore elements in ascending order by always choosing the smallest unexplored element adjacent to already processed elements.

Binary Search on Values: Instead of searching positions, we search the value space. For any candidate value, we can efficiently count how many elements are ≤ that value using the matrix’s sorted properties.

When to Apply: Use heap approach when k is small relative to matrix size. Use binary search when k is large or when you need to handle multiple queries. This pattern applies to any problem involving finding order statistics in structured data where you can efficiently count elements relative to a threshold.