LeetCode-Kth Smallest Element in a Sorted Matrix

因为矩阵中某一个元素下方和右边的数字比它大，所以在搜索的时候可以使用类似 Prim 算法的方式，在搜索的结果边缘进行搜索，然而算法超时。

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int m=matrix.size();
        int n=matrix[0].size();
        set<pair<int,int>> bag;
        pair<int,int> minPos=make_pair(0,0);
        bag.insert(minPos);
        for(int i=2;i<=k;i++){
            int minNum=INT_MAX;
             for(auto it=bag.begin();it!=bag.end();it++){
                pair<int,int> tPos=*(it);
                pair<int,int> newPos1=make_pair(tPos.first+1,tPos.second);
                pair<int,int> newPos2=make_pair(tPos.first,tPos.second+1);
                if((bag.find(newPos1)!=bag.end() || tPos.first==m-1) &&
                  (bag.find(newPos2)!=bag.end() || tPos.second==n-1)){
                    bag.erase(it);
                    continue;
                }
                if(bag.find(newPos1)==bag.end() && (*it).first<m-1 && matrix[newPos1.first][newPos1.second]<minNum){
                    minPos=newPos1;
                    minNum=matrix[newPos1.first][newPos1.second];
                }
                if(bag.find(newPos2)==bag.end() && (*it).second<n-1 && matrix[newPos2.first][newPos2.second]<minNum){
                    minPos=newPos2;
                    minNum=matrix[newPos2.first][newPos2.second];
                }
            }
            bag.insert(minPos);
        }
        return matrix[minPos.first][minPos.second];
    }
};

放弃复杂的尝试，使用优先级队列（堆）来实现也是一个比较方便的办法（可以通过剪枝优化一些时间复杂度）。但是这里几乎没有好好利用矩阵的性质。

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

发现 Discuss 中的二分法解法也很妙。

这里的二分法和常用在排序数组中的二分法有些不同，这里的二分是取最小和最大的数字（分别位于矩阵的左上角和右下角），然后将这两个数字进行二分，查找矩阵中小于该数字的元素的个数，如果元素的个数小于 k，那就说明要找的第 k 个元素在二分数字的右边，否则它就在二分数字的前面。这样就可以进一步缩小扫描的区间，最终将区域定位在要找的数字上。

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int m=matrix.size(),n=matrix[0].size();
        int start=matrix[0][0];
        int end=matrix[m-1][n-1];
        while(end>=start){
            int mid=start+(end-start)/2;
            int count=0;
            for(int i=0;i<m;i++){
                for(int j=0;j<n;j++){
                    if(matrix[i][j]<=mid) count++;
                    else break;
                }
            }
            if(count<k) start=mid+1;
            else end=mid-1;
        }
        return start;
    }
};

最后，还能使用Leetcode 240. Search a 2D Matrix II中的思路（从左下角开始比较）优化查找 mid 的时间复杂度（优化到对数级别）。

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int m=matrix.size(),n=matrix[0].size();
        int start=matrix[0][0];
        int end=matrix[m-1][n-1];
        while(end>=start){
            int mid=start+(end-start)/2;
            int count=0;
            int i=m-1,j=0;
            while(i>=0 && j<n){
                if(matrix[i][j]<=mid){
                    j++;
                    count+=i+1;
                }else{
                    i--;
                }
            }
            if(count<k) start=mid+1;
            else end=mid-1;
        }
        return start;
    }
};