【leetcode】378. 有序矩陣中第K小的元素（kth-smallest-element-in-a-sorted-matrix）（堆）[中等]

鏈接

https://leetcode-cn.com/problems/kth-smallest-element-in-a-sorted-matrix/

耗時

解題：~3 h
題解：16 min

題意

給定一個 n x n 矩陣，其中每行和每列元素均按升序排序，找到矩陣中第 k 小的元素。

思路

首先可以斷定最左上角的元素一定是最小的，並且對於每個元素都存在：$a_{i,j}<a_{i+1,j}$ 和 $a_{i,j}<a_{i,j+1}$。所以我想可以從最左上角的元素開始擴展，將向右邊和下邊擴展得到的每個元素放入優先隊列，這樣順序出隊的第 k 個元素即爲矩陣中第 k 小的元素。

時間複雜度：$O(klogk)$，最壞是$O(n^2logn)$

AC代碼

class Solution {
public:
    struct node {
        int val;
        int i, j;
        bool operator < (const node &d) const {
            return val > d.val;
        }
    };
    
    bool check(int x, int y, int n) {
        if((x >= 0 && x < n) && (y >= 0 && y < n)) 
            return true;
        return false;
    }
    
    void pq_push(int x, int y, int val, priority_queue<node> &pq) {
        node tmp;
        tmp.val = val;
        tmp.i = x;
        tmp.j = y;
        pq.push(tmp);

    }
    
    int vis[1100][1100] = {0};
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int n = matrix.size();
        priority_queue<node> pq;
        pq.push((node){matrix[0][0], 0, 0});
        vis[0][0] = 1;
        int cnt = 0, ans = -1;
        while(cnt < k) {
            node now = pq.top();
            pq.pop();
            cnt++;
            if(cnt == k) {
                ans = now.val;
                break;
            }
            
            if(check(now.i+1, now.j, n) && vis[now.i+1][now.j] == 0) {
                pq_push(now.i+1, now.j, matrix[now.i+1][now.j], pq);
                vis[now.i+1][now.j] = 1;
            }
            
            if(check(now.i, now.j+1, n) && vis[now.i][now.j+1] == 0) {
                pq_push(now.i, now.j+1, matrix[now.i][now.j+1], pq);
                vis[now.i][now.j+1] = 1;
            }
            
        }
        return ans;
    }
};