Task

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

Example:

1
2
3
4
5
6
7
8

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

return 13.

Note: You may assume k is always valid, 1 ≤ k ≤ n2.

Solution

  • 矩阵左上角和右下角分别是最小值和最大值，在最小值和最大值之间二分搜索出第k大的值
  • 每次得到二分的中间值m时都对整个矩阵统计小于m的元素数量：
    • 若数量小于k则m偏小，改变二分搜索左边界
    • 若数量大于k则m偏大，改变二分搜索右边界
  • 在每一行中统计小于m的元素数量时，由于每一行是排好序的，故也用二分搜索来统计。C++标准库函数upper_bound实现了这一功能，它的原理也是二分查找
  • 统计时注意到一个细节：由于每一列也是排好序的，故当前行中小于m的元素数量不会超过上一行的数量，可以此来缩小每一行统计的范围
  • 时间复杂度：二分次数O(log(max-min))，每次二分时遍历n行，遍历每一行时upper_bound二分查找的复杂度是O(logn)，故总时间复杂度是O(nlogn*log(max-min))
  • 空间复杂度：未引入额外变量，O(1)
  • 实现：
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

// Runtime: 56 ms, faster than 93.07% of C++ online submissions for Kth Smallest Element in a Sorted Matrix.
// Memory Usage: 13.2 MB, less than 64.10% of C++ online submissions for Kth Smallest Element in a Sorted Matrix.
class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        const int n=matrix.size();
        //对最小值和最大值之间的范围做二分搜索
        int l=matrix[0][0],r=matrix[n-1][n-1]+1;
        while(l<r){
            int m=l+(r-l)/2;
            int total_less=0,last_less=n;
            for(const auto &row:matrix){
                //使用upper_bound统计每一行中小于m的元素数量
                //同时更新last_less，统计下一行时只需统计last_less个元素，不需统计整行
                last_less=distance(row.begin(),upper_bound(row.begin(),row.begin()+last_less,m));
                //统计小于m的元素总数
                total_less+=last_less;
            }
            if(total_less<k)
                l=m+1;
            else
                r=m;
        }
        return l;
    }
};