378. 有序矩阵中第K小的元素

Q:

A:

//O(NK)
class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int m=matrix.size();
        int cnt=0;
        vector<int> indexes(m,0);
		int min_index=0,min_val=INT32_MAX;
        while(cnt<k){
			min_val=INT32_MAX;
			for(int i=0;i<m;++i){
				if (indexes[i]<m and matrix[i][indexes[i]]<min_val){
					min_val=matrix[i][indexes[i]];
					min_index=i;
				}
			}
			indexes[min_index]+=1;
			cnt++;
        }
		return min_val;
    }
};

最小堆，n路归并k次

//O(KlogN)
struct point
{
	int x,y,val;
	point():x(0),y(0),val(0){}
	point(int a,int b,int c):x(a),y(b),val(c){}
};
bool cmp(point& a,point& b){
	return a.val>b.val;
}
class Solution {
public:
	
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int m=matrix.size();
        int cnt=0;
        vector<point> vec;
		for(int i=0;i<m;i++){
			vec.push_back(point(i,0,matrix[i][0]));
		}
		make_heap(vec.begin(),vec.end(),cmp);
		int res=0;
        while(cnt<k){
			res=vec[0].val;
			pop_heap(vec.begin(),vec.end(),cmp);
			vec.back().y++;
			if (vec.back().y>=m){
				vec.pop_back();
			}
			else{
				point p=vec.back();
				vec.back().val=matrix[p.x][p.y];
			}
            if (vec.empty()){
                return res;
            }
			push_heap(vec.begin(),vec.end(),cmp);
            cnt++;
        }
		return res;
    }
};

二分，不断二分矩阵直到分出前k个元素。

//O(NlogN)
class Solution {
public:
    // O(nlogn)
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int lo = matrix[0][0], hi = matrix.back().back();
        // O(log(max_val - min_val)) = O(1)，32 位整型最多 32 次
        while (lo < hi) {
            int mid = lo + (hi - lo >> 1);
            int count = 0;
            // O(n)
            for (int i = 0; i < matrix.size(); ++i) {
                // O(logn)
                count += search(matrix[i], mid); // <= mid 的元素个数
            }
            if (k <= count) {
                hi = mid;
            } else {
                lo = mid + 1;
            }
        }
        return lo;
    }
    
    // 标准 upperbound 模板。查找第一个 > target 的元素的位置。
    // 也就是 <= target 的元素的个数。
    int search(vector<int>& row, int target) {
        int lo = 0, hi = row.size();
        while (lo < hi) {
            int mid = lo + (hi - lo >> 1);
            if (row[mid] > target) {
                hi = mid;
            } else {
                lo = mid + 1;
            }
        }
        return lo;
    }
};