[LeetCode] 378. Kth Smallest Element in a Sorted Matrix

题目描述

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

Note that it is the kth smallest element in the sorted order, not the kth distinct element.

Example:

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

return 13.

Note:
You may assume k is always valid, 1 ≤ k ≤ n2.
输入是一个二维矩阵，二维矩阵的每行每列都有序，要求求出矩阵所有数值的第k 小的数。
解题思路是使用一个优先队列(堆)存储矩阵的数值，优先队列pop (k - 1) 次后，位于堆顶的值就是第k 小的数值。由于二维矩阵的每行每列都有序，因此我们使用矩阵的第一行或者第一列初始化堆，每pop 一个节点，将节点下侧(右侧)的数入堆。
需要注意的是 c++ 中提供的priority_queue 默认是大根堆，需要重载节点的比较运算符实现小根堆的功能。

C++ 实现

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        const size_t m = matrix.size();
        if (matrix.empty()) return -1;
        const size_t n = matrix[0].size();

        struct Node
        {
            size_t x;
            size_t y;
            int val;

            bool operator < (const Node& other) const
            {
                return val >= other.val;
            }
        };

        priority_queue<Node> pq;
        for (size_t j = 0; j < n; j++)
        {
            pq.push({0, j, matrix[0][j]});
        }

        for (int i = 0; i < k - 1; i++)
        {
            const Node& top = pq.top();
            const size_t x = top.x;
            const size_t y = top.y;
            pq.pop();
            if (x == m - 1)
            {
                continue;
            }
            pq.push({x + 1, y, matrix[x + 1][y]});
        }
        return pq.top().val;
    }
};