378. Kth Smallest Element in a Sorted Matrix Medium

Given an n x n matrix where each of the rows and columns 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.

You must find a solution with a memory complexity better than O(n^2).

Example 1:
Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
Output: 13
Explanation: The elements in the matrix are [1,5,9,10,11,12,13,13,15], and the 8th smallest number is 13

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

Approach

Input: An n x n matrix

Output: Find the kth smallest element in the matrix

Method Analysis

• We can use a Min Heap to reduce complexity.

Specific Walkthrough

1. Initialize the heap:
   • Push the first element of each row (value, row, col) into the heap.
   • For example, the first element of row 0 is (matrix[0][0], 0, 0).
   • The storage structure in the heap is [value, row, col].
2. Pop k times repeatedly:
   • Every time pop the minimum value (val, r, c) from the heap;
   • If this row has a next element (i.e. c + 1 < n), push (matrix[r][c + 1], r, c + 1) into the heap;
   • Repeat this process k times, the element popped on the kth time is the answer.

Complexity Analysis

• The heap holds at most n elements, each pop/insert operation takes log n.
• Total Time Complexity: O(k log n)
• Space Complexity: O(n)

Implementation

python

import heapq
from typing import List

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        """
        Idea: Min Heap — Multi-way Merge Approach
        Every row is in ascending order. Treating each row as a sorted array,
        we can leverage a min heap to iteratively take the minimum available overall value.
        """

        n = len(matrix)
        heap = []  # Define a min heap, elements format: (value, row, col)

        # 1️⃣ Initialization: Push the first element of each row into the heap
        for i in range(n):
            # Push (matrix[i][0], i, 0) into the heap
            # Meaning: value = matrix[i][0], from row i, column 0
            heapq.heappush(heap, (matrix[i][0], i, 0))

        # 2️⃣ Pop the minimum value k-1 times, the kth minimum will be the answer
        for _ in range(k - 1):
            val, r, c = heapq.heappop(heap)  # Pop the current top array (minimum)

            # If the row corresponding to this element has a next element, push it into the heap
            if c + 1 < n:
                heapq.heappush(heap, (matrix[r][c + 1], r, c + 1))
        
        # 3️⃣ Now the top of the heap houses identically the kth minimum element requested
        return heap[0][0]

javascript

class MyMinHeap {
  constructor() {
    this.data = [];
  }
  size() { return this.data.length; }
  peek() { return this.data[0]; }
  push(val) {
    this.data.push(val);
    this._heapifyUp();
  }
  pop() {
    if (this.size() === 0) return null;
    const top = this.data[0];
    const end = this.data.pop();
    if (this.size() > 0) {
      this.data[0] = end;
      this._heapifyDown();
    }
    return top;
  }
  _heapifyUp() {
    let i = this.size() - 1;
    while (i > 0) {
      const p = Math.floor((i - 1) / 2);
      if (this.data[p][0] <= this.data[i][0]) break;
      [this.data[p], this.data[i]] = [this.data[i], this.data[p]];
      i = p;
    }
  }
  _heapifyDown() {
    let i = 0;
    const n = this.size();
    while (true) {
      const l = 2 * i + 1, r = 2 * i + 2;
      let smallest = i;
      if (l < n && this.data[l][0] < this.data[smallest][0]) smallest = l;
      if (r < n && this.data[r][0] < this.data[smallest][0]) smallest = r;
      if (smallest === i) break;
      [this.data[i], this.data[smallest]] = [this.data[smallest], this.data[i]];
      i = smallest;
    }
  }
}

/**
 * @param {number[][]} matrix
 * @param {number} k
 * @return {number}
 */
var kthSmallest = function(matrix, k) {
  const n = matrix.length;
  const heap = new MyMinHeap();

  // Initialization: The first element of every row gets pushed into the heap
  for (let i = 0; i < n; i++) {
    heap.push([matrix[i][0], i, 0]);
  }

  let result = null;
  // Pop iterations k times
  for (let i = 0; i < k; i++) {
    const [val, r, c] = heap.pop();
    result = val;
    // Pushing subsequent next column element on popped row safely boundaries allowing
    if (c + 1 < n) heap.push([matrix[r][c + 1], r, c + 1]);
  }

  return result;
};

Complexity Analysis

• Time Complexity: O(k log n)
• Space Complexity: O(n)

Links

378. Kth Smallest Element in a Sorted Matrix (English)

378. 有序矩阵中第 K 小的元素 (Chinese)