有序数组中第K小的数字

Posted 2020-12-14 realzhaijiayu

思路

首先 想到的就是暴力算法，将二维数组拆成一维数组，排序(或者直接放进优先队列)最后竟然过了。。。

因为矩阵在行列方向上都是升序排列的，在有序数组中搜索一个数字常用二分法，所以可以考虑变种的二分法。

二维数组从左下向右上搜索，每次搜索检查一下当前搜遍历到数字范围有没有超过k。

代码

暴力：

import java.util.Arrays;

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        int[] nums = new int[n*n];
        int index = 0;
        for(int i = 0; i < n; i++){
            for(int j = 0; j < n; j++){
                nums[index++] = matrix[i][j];
            }
        }

        Arrays.sort(nums);
        return nums[k-1];
    }
}

二分法

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;
    }

    /**
     * 检查小于等于mid的数字数量是否大于等于k
     * num表示小于等于mid的数字数量
     * @param matrix
     * @param mid
     * @param k
     * @param n
     * @return
     */
    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;
    }
}