378. 有序矩阵中第 K 小的元素

Posted 2022-04-13 hequnwang10

一、题目描述

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

你必须找到一个内存复杂度优于 O(n^2) 的解决方案。

示例 1：

输入：matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
输出：13
解释：矩阵中的元素为 [1,5,9,10,11,12,13,13,15]，第 8 小元素是 13

示例 2：

输入：matrix = [[-5]], k = 1
输出：-5

二、解题

优先队列+小顶堆

这里使用优先队列和小顶堆，将每一行的第一个元素入队，最小值一定在队列中，然后出队，将最小值弹出，将最小值所在行的下一个元素入队，重新将该队列排序，然后在出队，第二小的元素出队，所以循环k次，出队的元素就是第K小的元素。

class Solution 
    public int kthSmallest(int[][] matrix, int k) 
        //m=n 是一个n*n的矩阵
        int m = matrix.length,n = matrix[0].length;
        //使用一个优先队列 小顶堆来维护，弹出k次就是第k小元素
        PriorityQueue<int[]> pq = new PriorityQueue<>(new Comparator<int[]>()
            public int compare(int[] a,int[] b)
                return a[0]-b[0];
            
        );
        for(int i = 0;i < m;i++)
            //(i,0)是在记录候选人在矩阵中的位置
            pq.add(new int[]matrix[i][0],i,0);
            
        
        //一共弹k次：这里k-1次，return的时候1次
        for(int i = 0;i<k-1;i++)
            // System.out.println("pq: "+ pq.size());
            //弹出最小值
            int[] cur = pq.poll();
            // System.out.println("cur[0]: "+ cur[0]+"cur[1]:"+cur[1]+"cur[2]:"+cur[2]);
            //如果这一行还没被弹完  加入这一行的下一个候选人
            if(cur[2] != n-1)
                pq.add(new int[]matrix[cur[1]][cur[2]+1],cur[1],cur[2]+1);
            
        
        // return的时候1次
        return pq.poll()[0];
    

输出：

pq: 3
cur[0]: 1cur[1]:0cur[2]:0
pq: 3
cur[0]: 5cur[1]:0cur[2]:1
pq: 3
cur[0]: 9cur[1]:0cur[2]:2
pq: 2
cur[0]: 10cur[1]:1cur[2]:0
pq: 2
cur[0]: 11cur[1]:1cur[2]:1
pq: 2
cur[0]: 12cur[1]:2cur[2]:0
pq: 2
cur[0]: 13cur[1]:1cur[2]:2

二分查找

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

class Solution 
    public int kthSmallest(int[][] matrix, int k) 
        //m=n 是一个n*n的矩阵
        int m = matrix.length,n = matrix[0].length;
        //二分查找
        int left = matrix[0][0] ,right = matrix[m-1][n-1];
        while(left <= right)
            int mid = left+(right-left)/2;
            int cnt = 0;
            for(int i = 0;i<m;i++)
                for(int j = 0;j<n && matrix[i][j]<=mid;j++)
                    cnt++;
                
            
            if(cnt < k)
                left = mid + 1;
            else
                right = mid - 1;
            
        
        return left;
    

class Solution 
    public int kthSmallest(int[][] matrix, int k) 
        //二分查找
       int n = matrix.length;
        int left = matrix[0][0];
        int right = matrix[n - 1][n - 1];
        while (left < right) 
            int mid = left + ((right - left) >> 1);
            if (check(matrix, mid, k, n)) 
                right = mid;
             else 
                left = mid + 1;
            
        
        return left;
    

    public boolean check(int[][] matrix, int mid, int k, int n) 
        int i = n - 1;
        int j = 0;
        int num = 0;
        while (i >= 0 && j < n) 
            if (matrix[i][j] <= mid) 
                num += i + 1;
                j++;
             else 
                i--;
            
        
        return num >= k;