有序矩阵中第K小的元素

题目

题目链接：有序矩阵中第K小的元素

给定一个 n x n 矩阵，其中每行和每列元素均按升序排序，找到矩阵中第 k 小的元素。

请注意，它是排序后的第 k 小元素，而不是第 k 个元素。

输入描述

第一行为 k 的值和矩阵的 n 的值
后续为 n * n 矩阵的数字，以空格分割

输出描述

矩阵中第k小的元素

样例输入

8 3
1 5 9
10 11 13
12 13 15

样例输出

13

代码

import java.util.Scanner;

/**
 * 思路：
 *      如果直接用排序的话，时间复杂度为 O(2*n*n*logn)
 *
 *      因为这个二维数组每一行每一列都是有序的，所以可以通过二分来进行优化
 */
public class Main {
    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);
        int k = scanner.nextInt(), n = scanner.nextInt();
        int data[][] = new int[n][n];
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                data[i][j] = scanner.nextInt();
            }
        }
        System.out.println(findKthNum(data, k));
    }

    // 二分套二分，时间复杂度 O(n*logn*logn)
    public static int findKthNum(int[][] matrix, int k) {
        int row = matrix.length;
        int col = matrix[0].length;
        int low = matrix[0][0];
        int high = matrix[row - 1][col - 1];
        while (low <= high) {
            int count = 0;
            int mid = low + ((high - low) >> 1);
            for (int[] item : matrix) {
                count += binarySearchCount(item, mid);
            }
            if (count < k) {
                low = mid + 1;
            } else {
                high = mid - 1;
            }
        }
        // 返回的 low 表示满足 count >= k 的最小的数，这个数就是要找的第k个数
        return low;
    }

    // 返回数组 data 中个小于等于 k 的数的个数
    public static int binarySearchCount(int[] data, int k) {
        int begin = 0, end = data.length - 1;
        while (begin <= end) {
            int mid = (begin + end) >> 1;
            if (data[mid] <= k) {
                begin = mid + 1;
            } else {
                end = mid - 1;
            }
        }
        return begin;
    }
}