矩阵求前k大数

问题描述

最近看到一个面试题，给定一个m*n的矩阵和正整数k，其中行递增有序，列无序，求出前k大的数。这题目要是搁在平时自己刷的话，暴力算法也能解决，但是这是在面试过程中，因此一般面试官要求使用尽可能的优的算法来解决。暴力算法的话比较容易想到，把所有的元素排序一遍，那么就可以得到前k大的值，时间复杂度为O(mnlogmn)，或者多路归并，时间复杂度大致为O(m^2n)，这部分暴力解法因人而异，所以上述的时间复杂度略有出入也是正常的。但是如果在面试的时候说这几种方法显然面试官是不会满意的。

问题分析

我们可以看到它是要求输出前k大的数，那显然没有必要遍历所有的元素，而且行是有序的，可以利用好这一点，我们也可以像多路归并那样，只不过我们是从后往前遍历，这一做法是为了利用行有序的特点，只需遍历k次，每次输出当前最大的值就可以了。这一想法既可以引入额外数据结构也可以不引入额外数据结构来实现。

算法思路

普通数组实现

思路设计

因为行是递增有序的，因此最后一列的元素中最大的必然是当前最大的，选取第一大元素之后，列索引向左挪动，继续上述操作k次即可。代码如下：

    
/**
 * @Author:zxp
 * @Description:m*n的矩阵，行有序，列无序，找出前k大的值
 * @Date:10:49 2024/5/7
 */
   public List<Integer> getKNums(int[][] nums,int k){
        int m=nums.length,n=nums[0].length;
        int[] column = new int[m];
        Arrays.fill(column,n-1);//初始化数组，索引指向最后一列。
        List<Integer> result=new ArrayList<>();
        for(int i=0;i<k;i++){
            int maxRow=0,max=Integer.MIN_VALUE;
            for(int j=0;j<m;j++){//选当前维护数组中的最大元素并记录。
                if(column[j]>=0&&nums[j][column[j]]>max){
                    max=nums[j][column[j]];
                    maxRow=j;
                }
            }
            result.add(max);
            column[maxRow]--;
        }
        return result;
    }

可以看到这里我们维护了一个列数组，长度为m也就是数组的行的大小，每一个位置的值代表当前行的列索引，每次遍历比较维护数组中对应的值，最大的即为当前最大的，k次之后方法结束，结果集中就是最终的前k大的数。

时空复杂度

先来看时间复杂度，初始化数组，索引指向最后一列，这部分时间复杂度为O(m)，然后执行k次每次比较最大值，这部分时间复杂度为O(km)，所以整体时间复杂度为O((k+1)m).至于空间复杂度方面，此处额外引入了一个列索引数组，长度为m，所以空间复杂度为O(m).

优先级队列优化

思路设计

可以看到上面的普通数组实现有一个弊端：每次执行的时候都需要比较维护数组中的最大值，这部分时间复杂度为O(m)，有没有一种数据结构可以维护若干元素种的最大值或者最小值呢？优先级队列就登场了，它的底层数据结构是一个大根堆或者小根堆，取决于队首向队尾是递增还是递减。这边我们维护一个优先级队列，大小最多为m，并且队首向队尾是递减的，也就是队首的元素最大，所以底层是大根堆实现。每次从队首弹出一个元素，这个元素就是当前最大的，执行k次之后方法结束，调整这个大小为m的大根堆的时间复杂度为O(logm)。代码如下：

/**
 * @Author:zxp
 * @Description:m*n的矩阵，行有序，列无序，找出前k大的值
 * @Date:10:49 2024/5/7
 */
    public List<Integer> getKNumsByPri(int[][] nums,int k){
        PriorityQueue<int[]> priorityQueue=new PriorityQueue<>(new Comparator<int[]>(){
            @Override
            public int compare(int[] o1,int[] o2){
                return o2[2]-o1[2];
            }
        });//优先级队列，泛型为长度为3的数组，行、列、值。
        int m=nums.length;
        int n=nums[0].length;
        List<Integer> result=new ArrayList<>();
        for(int i=0;i<m;i++)
            priorityQueue.offer(new int[]{i,n-1,nums[i][n-1]});//队列初始化。
        for(int i=0;i<k;i++){
            int[] poll = priorityQueue.poll();
            result.add(poll[2]);
            if(poll[1]>0)
                priorityQueue.offer(new int[]{poll[0],poll[1]--,nums[poll[0]][poll[1]--]});
        }
        return result;
    }

可以看到这边维护了一个优先级队列，泛型为长度为3的数组，数组内容为行、列、值。从大根堆实现的优先级队列中每次弹出的数组中的值为当前队列中最大的，由于矩阵特性，它也是全局里前k大数的一部分。

时空复杂度

先来分析时间复杂度，初始化队列时间复杂度为O(m)，然后是执行k次每次弹出队列中队首元素并获取当前队列的最大值，调整大根堆时间复杂度为O(logm)，执行k次，所以这部分的时间复杂度为O(klogm)，所以整体时间复杂为O(m+klogm)。再来说一下空间复杂度，这边引入了一个大小最大为m的优先级队列，大根堆实现，因此空间复杂度为O(m)。

总结

求数组或者矩阵中的前k大元素的题目在面试中经常出现，它们的变体有很多种，本文所涉及到只是其中的一种形式。但是像取前k大元素的这种需求，一般会考虑到用堆来实现，大根堆还是小根堆看具体要求，其实优先级队列底层也是堆来实现的，所以最终还是回到了堆的灵活使用的问题上了，本人之前还探讨过数组取第k大元素的各种做法，那里本人也使用到了堆的思想，包括手撕堆以及引入优先级队列，还有快速排序的求解思路，感兴趣的小伙伴可以点击这个链接：求第k大数经典排序算法研究。本题整体来讲做出来并不难，如果要想到更优的做法，对堆的思想比较熟悉的话，会比较容易。当然本人相信本题的更优做法不止于此，如果有小伙伴有更好的做法，欢迎大家讨论。