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:

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.

这道题给我们一个二维数组,二维数组的行和列按升序排列,要求找出第k小的元素,k的范围是1 ≤ k ≤ n2.

对于寻找数组中第k小的元素,最经典的方法是使用堆这种数据结构,我们首先要建立一个大顶堆,然后让把数组中的元素依次插入堆中,当堆中元素数目超过k时,就删除堆顶元素,直到所有元素都遍历完后,返回堆顶元素的值即可。

这种方法根据大顶堆的性质很容易想到,时间复杂度为O(n^2 * logk),k是堆的大小,效率太低,而且没有用上题目中元素按行按列递增的性质。

代码如下:

/*solution 1,use heap,144ms,O(n^2 *  logk),k is the size of heap*/
class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        priority_queue<int> heap;
        for(int i=0;i<matrix.size();i++){
            for(int j=0;j<matrix[i].size();j++){
                heap.emplace(matrix[i][j]);
                if(heap.size()>k)
                    heap.pop();
            }
        }
        return heap.top();
    }
};

于是考虑能不能用二分查找对算法进行优化,考虑到元素按行按列递增,那么左上角必然是最小元素,右下角必然是最大元素,那么我们计算最小与最大的均值作为二分查找中的mid,假设mid就是我们的数组中第k小的元素,然后以mid元素为准,使用upper_bound()函数来统计每行中小于mid元素的个数,这个个数与k作比较,有两种情况:

  1. 大于或等于k,说明选择的mid太大,需要调整到左半区,则high=mid-1;
  2. 小于k,说明我们选定的mid太小,需要调整到右半区,则low=mid+1;

那么最后要返回的元素是什么?因为猜想的解不一定在数组中,所以要不断的收缩直到找到在数组中的元素为止。

时间复杂度为O(nlogn),运行时间56ms,代码如下:

/*solution 2,use binary-search,O(nlogn),56ms*/
class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int low=matrix[0][0],n=matrix.size(),high=matrix[n-1][n-1],cnt=0,mid=low+((high-low)>>1);
        while(low<=high){
            mid=low+((high-low)>>1);
            cnt=0;
            for(int i=0;i<matrix.size();i++)
                cnt=cnt+upper_bound(matrix[i].begin(),matrix[i].end(),mid)-matrix[i].begin();
            if(cnt<k)
                low=mid+1;
            else
                high=mid-1;
        }
        return low;
    }
};

北邮教务系统评教脚本

Published on September 17, 2017

72.Edit distance

Published on September 17, 2017