題目網址：https://leetcode.cn/problems/kth-smallest-element-in-a-sorted-matrix/

題意：給一 n x n matrix, 其中每行、每列各自按照升序排序, 返回 matrix 中第 k 小的元素。

注意：它是排序後第 k 小的元素, 而不是第 k 個不同的元素

e.g. matrix = [[1,2], [2,3]], k = 3, 答案是 2

➔ 因為排序後 [1, 2, 2, 3] 中第 3 小的是 2, 重複數不能只算一次（不能看做 [1, 2, 3]）

設計 $O(n^2)$ time 的演算法

進階：設計 $O(1)$ space 的演算法

Solution 1：

想法：利用 Heap, 由於 matrix 是已排序好的, 所以 (i, j) 的值一定比 (i + 1, j) 和 (i, j + 1) 的值小, 演算法可參考此影片, 簡單來說分成以下步驟：

  1. 取出 min heap 的 top, 也就是當前最小
  2. insert top 的右方和下方到 heap 中 (同位置不能重複 insert)
  3. 從 heap 中取出 k - 1 個 top 後, 剩餘 heap 中的 top 即為第 k 小的數

min Heap 定義：任意 parent 的值一定比 children 的都小

class Coord {
public:
    int r, c, val;

    Coord(int i, int j, int val) : r(i), c(j), val(val) {}

    bool operator> (const Coord& a) const { return this->val > a.val; }
};

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        const int n = matrix.size();
        vector<vector<bool>> visited(n, vector<bool>(n, false)); // 避免重複 insert
        priority_queue<Coord, vector<Coord>, greater<Coord>> pq; // min heap
        pq.emplace(Coord(0, 0, matrix[0][0]));

        // pop 掉前 k - 1 小的數, 則 pq.top() 即為第 k 小的數
        for (int i = 0; i < k - 1; ++i) {
            const auto [r, c, val] = pq.top();
            pq.pop();

            // 將 top 下方 insert 到 pq 中
            if (r + 1 < n && !visited[r + 1][c]) {
                visited[r + 1][c] = true;
                pq.emplace(Coord(r + 1, c, matrix[r + 1][c]));
            }

            // 將 top 右方 insert 到 pq 中
            if (c + 1 < n && !visited[r][c + 1]) {
                visited[r][c + 1] = true;
                pq.emplace(Coord(r, c + 1, matrix[r][c + 1]));
            }
        }

        return pq.top().val;
    }
};
  • time：$O(k \cdot log(n))$
    • $O(k)$：for loop, 其中 worse case 為 $k = n^2$
    • $O(log(n))$：insertion of heap（因為 heap 是 balanced tree, swap() 最多為樹高 $log(n)$ 次）
  • space：$O(n^2)$ ➔ visited

Solution 2：

想法：Solution 1 的改進, 藍色數字代表 minHeap 中元素(初始為第一行的元素), 彈出 minHeap 最小的後, 將該元素右邊一格加入到 heap 中成為下一個候選人

class Coord {
public:
    int r, c, val;

    Coord(int i, int j, int val) : r(i), c(j), val(val) {}

    bool operator> (const Coord& a) const { return this->val > a.val; }
};

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        const int n = matrix.size();
        priority_queue<Coord, vector<Coord>, greater<Coord>> minHeap;

        // 將第一行元素加到 minHeap 中
        for (int r = 0; r < n; ++r) {
            minHeap.emplace(r, 0, matrix[r][0]);
        }

        for (int i = 0; i < k - 1; ++i) {
            const auto [r, c, val] = minHeap.top();
            minHeap.pop();

            // 如果 top 右邊一格還沒越界, 則將其加入到 heap 中
            if (c != n - 1) {
                minHeap.emplace(r, c + 1, matrix[r][c + 1]);
            }
        }

        return minHeap.top().val;
    }
};
  • time：$O(k \cdot log(n))$
    • $O(k)$：for loop, 其中 worse case 為 $k = n^2$
    • $O(log(n))$：insertion of heap（因為 heap 是 balanced tree, swap() 最多為樹高 $log(n)$ 次）
  • space：$O(n)$ ➔ minHeap 中最多 n 個元素

Solution 3：

想法：利用 Binary Search, 其中 left = matrix[0][0], right = matrix[n - 1][n - 1] + 1, 得出 mid 後, 計算 matrix 中 ≤ mid 的個數 cnt (可參考 240. Search a 2D Matrix II 的 Z 字形搜索), 希望找到一個 mid 使得 matrix 中剛好有 k 個數 ≤ mid

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        const int n = matrix.size();
        int left = matrix[0][0], right = matrix[n - 1][n - 1] + 1;

        while (left < right) {
            const int mid = left + (right - left) / 2;

            if (count_less_equal(n, matrix, mid) >= k) {
                right = mid;
            } else {
                left = mid + 1;
            }
        }

        return left; // left 為第一個有 k 個數 >= 自己
    }

private:
    // 計算 matrix 中 ≤ mid 的元素個數, 此處也可以改成用 binary search 實作
    int count_less_equal(const int n, const vector<vector<int>>& matrix, const int mid){
        int cnt = 0, row = 0, col = n - 1;

        while (row < n && col >= 0) {
            if (matrix[row][col] <= mid) {
                cnt += (col + 1);
                ++row;
            } else {
                --col;
            }
        }

        return cnt;
    }
};
  • time：$O(n \cdot log(r-l))$ ➔ r 代表 right, l 代表 left
    • $O(log(r-l))$：Binary Search 的次數
    • $O(n)$：計算 cnt 的時間 $O(n+n)$
  • space：$O(1)$ ➔ 只需常數空間