Kth smallest element in a row-column wise sorted matrix

Problem Statement

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

Example:

• 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 is 13.

Approach 1: Brute Force (Flatten and Sort)

Intuition

Copy all elements to a 1D array, sort it, and return the kth element.

Algorithm

1. Flatten the matrix into a 1D array
2. Sort the array
3. Return element at index k-1

import java.util.*;

public class Solution {
    public int kthSmallestBrute(int[][] matrix, int k) {
        int n = matrix.length;
        int[] arr = new int[n * n];
        int index = 0;
        
        for (int[] row : matrix) {
            for (int val : row) {
                arr[index++] = val;
            }
        }
        
        Arrays.sort(arr);
        return arr[k - 1];
    }
}

Complexity Analysis

• Time Complexity: O(n^2 * log(n^2)) = O(n^2 * log(n)) - Sorting n^2 elements
• Space Complexity: O(n^2) - Storing all elements

Approach 2: Min-Heap (Merge K Sorted Lists)

Intuition

Use a min-heap to merge n sorted rows. Extract k elements to find the kth smallest.

Algorithm

1. Initialize heap with first element of each row
2. Extract minimum k times
3. After each extraction, add next element from the same row

import java.util.*;

public class Solution {
    public int kthSmallestHeap(int[][] matrix, int k) {
        int n = matrix.length;
        
        // Min-heap: {value, row, col}
        PriorityQueue<int[]> minHeap = new PriorityQueue<>((a, b) -> a[0] - b[0]);
        
        // Initialize heap with first column
        for (int i = 0; i < n; i++) {
            minHeap.offer(new int[]{matrix[i][0], i, 0});
        }
        
        int result = 0;
        
        // Extract k elements
        for (int i = 0; i < k; i++) {
            int[] curr = minHeap.poll();
            result = curr[0];
            int row = curr[1], col = curr[2];
            
            // Add next element from same row
            if (col + 1 < n) {
                minHeap.offer(new int[]{matrix[row][col + 1], row, col + 1});
            }
        }
        
        return result;
    }
}

Complexity Analysis

• Time Complexity: O(k * log(n)) - k extractions from heap of size n
• Space Complexity: O(n) - Heap stores at most n elements

Approach 3: Binary Search on Value Range (Optimal)

Intuition

Binary search on the value range [min, max]. For each mid value, count how many elements are ≤ mid. Use this count to narrow down the search.

The kth smallest element is the smallest value such that count of elements ≤ value is ≥ k.

Algorithm

1. Set low = matrix[0][0], high = matrix[n-1][n-1]
2. Binary search:
   • For mid, count elements ≤ mid using staircase traversal
   • If count < k, search higher
   • Else search lower
3. Return low when low == high

public class Solution {
    private int countLessOrEqual(int[][] matrix, int target) {
        int n = matrix.length;
        int count = 0;
        int row = n - 1, col = 0;  // Start from bottom-left
        
        while (row >= 0 && col < n) {
            if (matrix[row][col] <= target) {
                count += row + 1;  // All elements in this column up to row
                col++;
            } else {
                row--;
            }
        }
        
        return count;
    }
    
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        int low = matrix[0][0], high = matrix[n-1][n-1];
        
        while (low < high) {
            int mid = low + (high - low) / 2;
            int count = countLessOrEqual(matrix, mid);
            
            if (count < k) {
                low = mid + 1;
            } else {
                high = mid;
            }
        }
        
        return low;
    }
}

Complexity Analysis

• Time Complexity: O(n * log(max - min))
  • Binary search: O(log(max - min)) iterations
  • Each count: O(n) using staircase
• Space Complexity: O(1) - Only using constant extra space

Why Binary Search Works

The key insight: if we search for value V and count of elements ≤ V is exactly k, V might not be in the matrix. But the algorithm converges to the smallest value with count ≥ k, which must be in the matrix.

Comparison of Approaches

| Approach | Time | Space | Best When | |----------|------|-------|-----------| | Brute Force | O(n² log n) | O(n²) | Simple implementation | | Min-Heap | O(k log n) | O(n) | k is small | | Binary Search | O(n log(max-min)) | O(1) | k is large, space critical |

Key Takeaways

1. Binary search on answer is optimal for sorted matrix queries
2. The staircase counting technique gives O(n) count in sorted matrix
3. Min-heap approach is better when k << n²
4. Binary search approach doesn't depend on k value
5. The answer must exist in matrix - binary search converges to it
6. This pattern applies to finding median, kth largest, etc.