[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 題目講解匯總(持續更新中...)