algoadvance

Leetcode 378. Kth Smallest Element in a Sorted Matrix

Problem Statement

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

Example:

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

Output: 13

Clarifying Questions

  1. Matrix dimensions and constraints:
    • What is the size range for n (the dimensions of the matrix)?
    • Constraints on values within the matrix and the range of k.
  2. Input properties:
    • Can there be duplicate elements in the matrix?
  3. Output requirements:
    • Should the function return the element itself or some indicator like an index?

Strategy

Approach:

  1. Min-Heap Method: Utilize a heap to keep track of the smallest elements efficiently.
    • Insert the first element of each row into the heap with its coordinates.
    • Pop the smallest element from the heap and push the next element of the same row into the heap.
    • Repeat the above process until we pop the kth element.

Steps:

  1. Initialize the Min-Heap: Insert the first element of each row into a min-heap.
  2. Pop and Push: Pop the smallest element from the heap and push the next element from the same row into the heap.
  3. Repeat: Continue the above step k times.
  4. Return the kth Element: The element popped on the kth iteration is the kth smallest element.

Code:

#include <vector>
#include <queue>
#include <tuple>

using namespace std;

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        // Min-heap to store elements as (value, row, col)
        auto comp = [](const tuple<int, int, int>& a, const tuple<int, int, int>& b) {
            return get<0>(a) > get<0>(b);
        };
        
        priority_queue<tuple<int, int, int>, vector<tuple<int, int, int>>, decltype(comp)> minHeap(comp);
        
        int n = matrix.size();
        
        // Initialize the heap with first element of each row.
        for (int r = 0; r < min(n, k); ++r) {
            minHeap.push({matrix[r][0], r, 0});
        }
        
        int value = 0;
        
        for (int i = 0; i < k; ++i) {
            auto [val, r, c] = minHeap.top();
            minHeap.pop();
            value = val;
            
            if (c + 1 < n) {
                // Push the next element in the row into the heap.
                minHeap.push({matrix[r][c + 1], r, c + 1});
            }
        }
        
        return value;
    }
};

Time Complexity

Hence, the overall time complexity is O(k log(min(k, n))).

This method ensures efficiency by leveraging the properties of a heap to manage a dynamic set of elements efficiently.

Cut your prep time in half and DOMINATE your interview with AlgoAdvance AI