Kth Smallest Element in Sorted matrix

97.Kth Smallest Element in a Sorted Matrix

Heap

Matrix

Binary Search

Problem Statement:

Given an n x n matrix where each of the rows and columns is sorted in ascending order, return the k^th smallest element in the matrix.

Algorithm:

We can solve this problem using a min-heap (priority queue) to efficiently find the kth smallest element:

1. Initialize a min-heap and insert the first element of each column (i.e., the elements of the first row).
2. While k is greater than zero:
• Extract the smallest element from the heap.
• If there is a next element in the same column, insert it into the heap.
• Decrement k by one.
3. The last extracted element is the kth smallest element.

Complexity:

Time: O(k log n), where n is the number of rows (or columns) in the matrix | Space: O(n)

Java Implementation:

// Classic Priority Queue approach
// Add all the elements from the first row
public int kthSmallest(int[][] matrix, int k) {
    int n = matrix.length;
    // Min-heap to store the elements
    PriorityQueue minHeap = new PriorityQueue<>((a, b) -> a.val - b.val);
    
    // Initialize the heap with the first element of each column
    for (int j = 0; j < matrix[0].length; j++)
        minHeap.offer(new Element(matrix[0][j], 0, j));
    
    // Extract the smallest elements k-1 times
    while (k-- > 1) {
        Element curr = minHeap.poll();
        if (curr.x < n - 1) {
            // If there is a next element in the same column, add it to the heap
            int nextVal = matrix[curr.x + 1][curr.y];
            minHeap.offer(new Element(nextVal, curr.x + 1, curr.y));
        }
    }
    
    // The kth smallest element
    return minHeap.poll().val; 
}

// Helper class to store the value and its position in the matrix
class Element {
    int val;
    int x; // Row index
    int y; // Column index
    
    public Element(int val, int x, int y) {
        this.val = val;
        this.x = x;
        this.y = y;
    }  
}
                

Python Implementation:

import heapq

def kthSmallest(matrix, k):
    n = len(matrix)
    minHeap = []
    
    # Initialize the heap with the first element of each column
    for j in range(len(matrix[0])):
        heapq.heappush(minHeap, (matrix[0][j], 0, j))
    
    # Extract the smallest elements k-1 times
    for _ in range(k - 1):
        val, x, y = heapq.heappop(minHeap)
        if x < n - 1:
            # If there is a next element in the same column, add it to the heap
            nextVal = matrix[x + 1][y]
            heapq.heappush(minHeap, (nextVal, x + 1, y))
    
    # The kth smallest element
    return heapq.heappop(minHeap)[0]
                

C++ Implementation:

#include 
#include 
using namespace std;

int kthSmallest(vector>& matrix, int k) {
    int n = matrix.size();
    // Min-heap to store the elements
    auto cmp = [](const Element& a, const Element& b) { return a.val > b.val; };
    priority_queue, decltype(cmp)> minHeap(cmp);
    
    // Initialize the heap with the first element of each column
    for (int j = 0; j < matrix[0].size(); j++)
        minHeap.emplace(matrix[0][j], 0, j);
    
    // Extract the smallest elements k-1 times
    for (int i = 0; i < k - 1; i++) {
        Element curr = minHeap.top();
        minHeap.pop();
        if (curr.x < n - 1) {
            // If there is a next element in the same column, add it to the heap
            int nextVal = matrix[curr.x + 1][curr.y];
            minHeap.emplace(nextVal, curr.x + 1, curr.y);
        }
    }
    
    // The kth smallest element
    return minHeap.top().val;
}

// Helper struct to store the value and its position in the matrix
struct Element {
    int val;
    int x; // Row index
    int y; // Column index
    
    Element(int val, int x, int y) : val(val), x(x), y(y) {}
};