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

時間 2019-11-30

標籤 leetcode kth smallest element 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.html

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

Example:github

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，大於的話就將首元素去掉，循環結束後咱們返回堆中的首元素即爲所求:less

解法一：函數

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小的數字。舉個例子來講吧，好比數組爲:post

[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爲最大值和最小值的差值，參見代碼以下：url

解法二：spa

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