【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^{2}logn)$

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;
    }
};