[LeetCode] 378、有序矩阵中第K小的元素

题目描述

给定一个 n x n 矩阵，其中每行和每列元素均按升序排序，找到矩阵中第k小的元素。
请注意，它是排序后的第k小元素，而不是第k个元素。

示例：

matrix = [
   [ 1,  5,  9],
   [10, 11, 13],
   [12, 13, 15]
],
k = 8,

返回 13。

解题思路

和《剑指offer》第4题类似，但是不一样。其实这个题的“二分法”的思路跟《剑指offer》的面试题3的题目二思路更像。

• 二分法

  • 找出二维矩阵中最小的数left，最大的数right，那么第k小的数必定在left~right之间。
  • mid=(left+right) / 2；在二维矩阵中寻找小于等于mid的元素个数count。
  • 若这个count小于k，表明第k小的数在右半部分且不包含mid，即left=mid+1, right=right，又保证了第k小的数在left~right之间。
  • 若这个count大于k，表明第k小的数在左半部分且可能包含mid，即left=left, right=mid，又保证了第k小的数在left~right之间。
  • 因为每次循环中都保证了第k小的数在left~right之间，当left==right时，第k小的数即被找出，等于right。

  注：这里的left mid right是数值，不是索引位置。

参考代码

typedef long long LL;
class Solution
{
public:
    int kthSmallest(vector<vector<int>>& matrix, int k){
        int rows = matrix.size();
        int cols = matrix[0].size();
        LL left = matrix[0][0];
        LL right = matrix[rows-1][cols-1];

        while(left < right){
            // 每次循环都保证第K小的数在left~right之间，当left==right，第k小的数就是left
            LL mid = (left + right) >> 1;
            // 找二维矩阵中<=mid的元素总个数
            LL cnt = findNotBiggerThanMid(matrix, rows, cols, mid);
            
            if(cnt < k)
                // 第k小的数在右半部分，且不包含mid
                left = mid + 1;
            else
                // 第k小的数在左半部分，可能包含mid
                right = mid;
        }
        
        return left;
    }
	
	// 根据有序矩阵的性质，优化查找过程
    int findNotBiggerThanMid(vector<vector<int>>& matrix, int rows, int cols, int mid){
    	// 从左下角开始找
        int i = rows - 1;
        int j = 0;
        int cnt = 0;
        while(i >= 0 && j < cols){
            if(matrix[i][j] <= mid){
                // 第j列有i+1个元素<=mid
                cnt += (i+1);
                j++;
            }
            else{
                // 第i行目前的数大于mid，需要继续在当前列往上找
                i--;
            }
        }

        return cnt;
    }
    
};

最大堆

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[0].size(); ++j) {
                pq.push(matrix[i][j]);
                if (pq.size() > k) pq.pop();
            }
        }
        return pq.top();
    }
};