【九章算法强化班】第k大

发表于 | 分类于 ， | 浏览 次
字数统计: 1,529 | 阅读时长 ≈ 7

求数组/矩阵的第k大元素

涉及leetcode题目

  1. 278. Kth Smallest Number In Matrix
  2. 373. Kth Smallest Sum In Two Sorted Arrays
  3. Kth Largest in N Arrays

278. Kth Smallest Number In Matrix

题目

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.

给定一个的矩阵，满足行递增和列递增，要求返回第k大的元素。

example

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

分析

方法1：堆

看到第k大问题，要想到用堆

因为矩阵满足行递增和列递增，所以矩阵中的每一个元素的右边和下边元素一定比这个元素大。可以从左上角开始将元素放入一个最小堆中，每次从最小堆中取出一个元素，将其右边和下边的元素放入最小堆，这样就可以保证直到取出第k次的元素就是矩阵中第k大的元素。

struct Node{
    int x;
    int y;
    int val;
    Node(int a,int b,int valin):x(a),y(b),val(valin){}
};
struct cmp{
    bool operator()(const Node &a,const Node &b)
    {
         return a.val>b.val;
    }
};
int kthSmallest(vector<vector<int>>& matrix, int k) {
    int rows = matrix.size();
    int cols = matrix[0].size();
    if(rows*cols<k){
        return -1;
    }
    vector<vector<bool> > isin(rows,vector<bool>(cols,0));//用于记录某元素是否已经入过堆
    priority_queue<Node,vector<Node>,cmp> minheap;
    int i = 0;
    int j = 0;
    Node nodetemp(i,j,matrix[i][j]);
    minheap.push(nodetemp);
    isin[i][j]=1;
    int kkk=0;
    while(kkk < k-1){
        int xtemp = minheap.top().x;
        int ytemp = minheap.top().y;
        int xadd = xtemp+1;
        int yadd = ytemp+1;
        if(xadd>=matrix.size()||ytemp>=matrix[0].size()||isin[xadd][ytemp]){
        }
        else{
            minheap.push(Node(xadd,ytemp,matrix[xadd][ytemp]));
            isin[xadd][ytemp]=1;
        }      if(yadd>=matrix.size()||xtemp>=matrix[0].size()||isin[xtemp][yadd]){
        }
        else{
            minheap.push(Node(xtemp,yadd,matrix[xtemp][yadd]));
            isin[xtemp][yadd]=1;
        }        
        minheap.pop();
        kkk++;
    }
    return minheap.top().val;
}

方法2：二分查找

利用堆的方法可以解决上面的问题，但是时间复杂度不是很好。

看了leetcode题解，有二分查找的方式更快。

先找到矩阵的最大值max和最小值min，即左上和右下元素，然后以此作为二叉搜索的左右两边，求出其中间值mid，遍历矩阵元素，求出比该中值小的元素的个数count，分一下三种情况讨论：

  • count<k，说明比mid小的元素个数不足k个，则要寻找的值比mid大，故mid++;
  • count>=k，说明比mid小或和mid值相等的元素个数多于k个，则要寻找的值<=mid，故令max = mid，继续查找;
  • 直到min和max回合，此时就找到了第k个元素。
int kthSmallest(vector<vector<int>>& matrix, int k) {
    int rows = matrix.size();
    int cols = matrix[0].size();
    if(rows*cols<k){
        return -1;
    }
    int min = matrix[0][0];
    int max = matrix[rows-1][cols-1];
    while(min < max){
        int mid = (min+max)/2;
        int count = 0;//matrix中记录比mid小的元素的个数
        for(int i = 0;i < rows;i++){
            for(int j = 0; j < cols;j++){
                if(matrix[i][j]<=mid){
                    count++;
                }
                else{
                    break;
                }
            }
        }
        if(count>=k){
            max = mid;
        }
        else{
            min = mid+1;
        }
    }
    return min;
}

373. Find K Pairs with Smallest Sums

题目

You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k.

Define a pair (u,v) which consists of one element from the first array and one element from the second array.

Find the k pairs (u1,v1),(u2,v2) …(uk,vk) with the smallest sums.

给定一个整数k和两个数组和，两个数组内部升序排列，分别从两个数组中选择一个元素和构成数对，返回和最小的前k个数对

Example 1:

Given nums1 = [1,7,11], nums2 = [2,4,6],  k = 3
  
Return: [1,2],[1,4],[1,6]
The first 3 pairs are returned from the sequence:
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

分析

这道题就是278的升级版，两个升序数组中各取一个元素做和，求最小的前k个，可以把两个升序数组看成是矩阵的两个维度：

1 7 11
2 2+1=3 2+7=9 2+11=13
4 4+1=5 4+7=11 4+11=15
6 6+1=7 6+7=13 6+11=17

显然，由元素对生成的矩阵和上题中具有相同的性质，同行同列都递增，所以只需要对上面的代码做轻微的调整即可：记录下元素对在原数组中的idx。

vector<pair<int,int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
    vector<pair<int, int>> results;
    if(nums1.size()==0||nums2.size()==0)
        return results;
    int len1 = nums1.size();
    int len2 = nums2.size();
    vector<vector<bool> > isin(len1,vector<bool>(len2,0));
    priority_queue<Node,vector<Node>,cmp> minheap;
    int i = 0;
    int j = 0;
    Node nodetemp(i,j,nums1[i]+nums2[j]);
    minheap.push(nodetemp);
    isin[i][j]=1;
    int numslen = len1*len2;
    for(int kkk = 0; kkk < k&&kkk<numslen;kkk++){
        int xtemp = minheap.top().x;
        int ytemp = minheap.top().y;
        results.push_back(make_pair(nums1[xtemp],nums2[ytemp]));
        int xadd = xtemp+1;
        int yadd = ytemp+1;
        if(xadd>=nums1.size()||ytemp>=nums2.size()||isin[xadd][ytemp]){
        }
        else{
            minheap.push(Node(xadd,ytemp,nums1[xadd]+nums2[ytemp]));
            isin[xadd][ytemp]=1;
        }
        if(xtemp>=nums1.size()||yadd>=nums2.size()||isin[xtemp][yadd]){
        }
        else{
            minheap.push(Node(xtemp,yadd,nums1[xtemp]+nums2[yadd]));
            isin[xtemp][yadd]=1;
        }        
        minheap.pop();
    }
    return results;
}

Kth Largest in N Arrays

题目

给定N个无序数组，从中找出第k大的元素

分析

看到数组先排序，然后再想思路。寻找第k大的元素，需要维护一个堆。

  1. 将N个数组中的最大值入堆
  2. 每次出堆一个元素，将该元素所在数组中的下一个入堆
  3. 循环k次，找到第k大的元素

依旧沿用上面题的思路，这次除元素大小外，需要记录的额外信息为元素所在的数组idx1，以及在该数组中的idx2。

int kthSmallest(vector<vector<int>>& matrix, int k) {
    int rows = matrix.size();
  	int cols = matrix[0].size();
    if(rows*cols<k){
        return -1;
    }
    priority_queue<Node,vector<Node>,cmp> minheap;
    //按行排序,每行第一个入堆
    for(inti= 0 ;i<rows;i++){
      sort(matrix[i].begin(),matrix[i].end());
      Node nodetemp(i,0,matrix[i][0]);
      minheap.push(nodetemp);
      isin[i][j]=1;
    }
    vector<vector<bool> > isin(rows,vector<bool>(cols,0));//用于记录某元素是否已经入过堆  
    int kkk=0;
    while(kkk < k-1){
        int xtemp = minheap.top().x;
      	int ytemp = minheap.top().y;
       	int yadd = ytemp++;
        if(yadd>=matrix[x].size()||isin[xtemp][yadd]){
        }
        else{
            minheap.push(Node(xtemp,yadd,matrix[xtemp][yadd]));
            isin[xadd][ytemp]=1;
        }
        minheap.pop();
        kkk++;
    }
    return minheap.top().val;
}

总结：

  • 见到需要维护集合的最小/最大值的时候要想到堆
  • 见到第k小，想到用堆维护候选集合，出堆k次
  • 见到数组要往排序上面想，先排序，然后再其他操作