[LeetCode] 378. Kth Smallest Element in a Sorted Matrix 有序矩阵中第K小的元素

 

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 ≤ n².

 

这道题让咱们求有序矩阵中第K小的元素，这道题的难点在于数组并非蛇形有序的，意思是当前行的最后一个元素并不必定会小于下一行的首元素，因此咱们并不能直接定位第K小的元素，因此只能另辟蹊径。先来看一种利用堆的方法，咱们使用一个最大堆，而后遍历数组每个元素，将其加入堆，根据最大堆的性质，大的元素会排到最前面，而后咱们看当前堆中的元素个数是否大于k，大于的话就将首元素去掉，循环结束后咱们返回堆中的首元素即为所求:

 

解法一：

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        priority_queue<int> q;
        for (int i = 0; i < matrix.size(); ++i) {
            for (int j = 0; j < matrix[i].size(); ++j) {
                q.emplace(matrix[i][j]);
                if (q.size() > k) q.pop();
            }
        }
        return q.top();
    }
};

 

这题咱们也能够用二分查找法来作，咱们因为是有序矩阵，那么左上角的数字必定是最小的，而右下角的数字必定是最大的，因此这个是咱们搜索的范围，而后咱们算出中间数字mid，因为矩阵中不一样行之间的元素并非严格有序的，因此咱们要在每一行都查找一下 mid，咱们使用 upper_bound，这个函数是查找第一个大于目标数的元素，若是目标数在比该行的尾元素大，则 upper_bound 返回该行元素的个数，若是目标数比该行首元素小，则 upper_bound 返回0, 咱们遍历完全部的行能够找出中间数是第几小的数，而后k比较，进行二分查找，left 和 right 最终会相等，而且会变成数组中第k小的数字。举个例子来讲吧，好比数组为:

[1 2
12 100]
k = 3
那么刚开始 left = 1, right = 100, mid = 50, 遍历完 cnt = 3，此时 right 更新为 50
此时 left = 1, right = 50, mid = 25, 遍历完以后 cnt = 3, 此时 right 更新为 25
此时 left = 1, right = 25, mid = 13, 遍历完以后 cnt = 3, 此时 right 更新为 13
此时 left = 1, right = 13, mid = 7, 遍历完以后 cnt = 2, 此时 left 更新为8
此时 left = 8, right = 13, mid = 10, 遍历完以后 cnt = 2, 此时 left 更新为 11
此时 left = 11, right = 12, mid = 11, 遍历完以后 cnt = 2, 此时 left 更新为 12
循环结束，left 和 right 均为 12，任意返回一个便可。

本解法的总体时间复杂度为 O(nlgn*lgX)，其中X为最大值和最小值的差值，参见代码以下：

 

解法二：

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int left = matrix[0][0], right = matrix.back().back();
        while (left < right) {
            int mid = left + (right - left) / 2, cnt = 0;
            for (int i = 0; i < matrix.size(); ++i) {
                cnt += upper_bound(matrix[i].begin(), matrix[i].end(), mid) - matrix[i].begin();
            }
            if (cnt < k) left = mid + 1;
            else right = mid;
        }
        return left;
    }
};

 

上面的解法还能够进一步优化到 O(nlgX)，其中X为最大值和最小值的差值，咱们并不用对每一行都作二分搜索法，咱们注意到每列也是有序的，咱们能够利用这个性质，从数组的左下角开始查找，若是比目标值小，咱们就向右移一位，并且咱们知道当前列的当前位置的上面全部的数字都小于目标值，那么 cnt += i+1，反之则向上移一位，这样咱们也能算出 cnt 的值。其他部分跟上面的方法相同，参见代码以下：

 

解法三：

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int left = matrix[0][0], right = matrix.back().back();
        while (left < right) {
            int mid = left + (right - left) / 2;
            int cnt = search_less_equal(matrix, mid);
            if (cnt < k) left = mid + 1;
            else right = mid;
        }
        return left;
    }
    int search_less_equal(vector<vector<int>>& matrix, int target) {
        int n = matrix.size(), i = n - 1, j = 0, res = 0;
        while (i >= 0 && j < n) {
            if (matrix[i][j] <= target) {
                res += i + 1;
                ++j;
            } else {
                --i;
            }
        }
        return res;
    }
};

 

Github 同步地址：

https://github.com/grandyang/leetcode/issues/378

 

相似题目：

Find K Pairs with Smallest Sums

Find K-th Smallest Pair Distance

Find K Closest Elements

Kth Smallest Number in Multiplication Table

K-th Smallest Prime Fraction

 

参考资料：

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

https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/discuss/85177/Java-1ms-nlog(max-min)-solution

https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/discuss/85222/C%2B%2B-priority-queue-solution-O(klogn)

https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/discuss/85182/My-solution-using-Binary-Search-in-C%2B%2B

 

LeetCode All in One 题目讲解汇总(持续更新中...)