378. Kth Smallest Element in a Sorted Matrix

Description

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

Intuition

Since the matrix is sorted in ascending order both row-wise and column-wise, the k-th smallest element can be efficiently found using a max-heap. The idea is to keep track of the k smallest elements seen so far. By using a max-heap, we can easily ensure that only the k smallest elements are retained.

Approach

1. Initialize a Max-Heap: Use a max-heap to keep track of the k smallest elements. In Python, a min-heap is used, but by storing negative values, it simulates a max-heap.
2. Iterate Through the Matrix: For each element in the matrix: - Add the element to the max-heap. - If the heap size exceeds k, remove the largest element from the heap to maintain only the k smallest elements.
3. Return the k-th Smallest Element: After processing all elements, the root of the max-heap will be the k-th smallest element.

Complexity

Time Complexity: O(m * n * log k)

• Inserting an element into the heap takes O(log k) time. Since we process each element of the matrix (m * n elements), the total time complexity is O(m * n * log k).

Space Complexity: O(k)

• The space complexity is determined by the heap, which stores at most k elements. Therefore, the space complexity is O(k).

Code

C++

class Solution {
public:
    int kthSmallest(vector>& matrix, int k) {
        priority_queue maxHeap; // Max-Heap for storing k smallest elements

        for (const auto& row : matrix) {
            for (int num : row) {
                maxHeap.push(num);
                if (maxHeap.size() > k) {
                    maxHeap.pop();
                }
            }
        }

        return maxHeap.top();
    }
};

Python

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

        for row in matrix:
            for num in row:
                heapq.heappush(max_heap, -num)  # Use negative values to simulate a max-heap
                if len(max_heap) > k:
                    heapq.heappop(max_heap)

        return -max_heap[0]

Java

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        PriorityQueue maxHeap = new PriorityQueue<>((a, b) -> b - a); // Max-Heap for k smallest elements

        for (int[] row : matrix) {
            for (int num : row) {
                maxHeap.offer(num);
                if (maxHeap.size() > k) {
                    maxHeap.poll();
                }
            }
        }

        return maxHeap.peek();
    }
}

JavaScript

var kthSmallest = function (matrix, k) {
    const heap = new MaxPriorityQueue()
    
    for (let i = 0; i < matrix.length; i++) {
        for (num of matrix[i]) {
            heap.enqueue(num)
            if (heap.size() > k) {
                heap.dequeue()
            }
        }
    }
    return heap.front().element

};