350，有序矩阵中第K小的元素

想了解更多数据结构以及算法题，可以关注微信公众号“数据结构和算法”，每天一题为你精彩解答。也可以扫描下面的二维码关注

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

示例:

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


返回 13。

答案：

public int kthSmallest(int[][] matrix, int k) {
    int n = matrix.length;
    PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> a[2] - b[2]);
    for (int j = 0; j < n; j++)
        pq.offer(new int[]{0, j, matrix[0][j]});
    for (int i = 0; i < k - 1; i++) {
        int[] cur = pq.poll();
        if (cur[0] == n - 1)
            continue;
        pq.offer(new int[]{cur[0] + 1, cur[1], matrix[cur[0] + 1][cur[1]]});
    }
    return pq.poll()[2];
}

解析：

PriorityQueue会对添加进去的数据进行排序，其实他就是一个堆，在这里他是个最小堆，也就是最顶端的元素是最小的（虽然他是数组结构，但数组的位置是有关联的），每添加一个元素他都会往上调整，删除的时候往下调整，并且他有两个最重要的函数一个是siftDown，一个是siftUp。上面的代码中它添加的是个数组，数组的最后一个元素是我们添加的值，前两个元素是这个值在矩阵中（x，y）的坐标。他是先把矩阵中第一行的元素全部添加进去，后面再进行k-1次循环。每次循环的时候都会把最小的给移除掉，然后把它紧挨着的下一个元素添加进去。因为是最小堆，所以当我们移除了k-1次的时候，这时堆的顶端就是第k小的元素。我们还以上面的例子来画个图加深一下理解

题中矩阵的每行每列都是排过序的，所以我们还可以想到另一种方法，使用二分法查找，之前介绍过二分法查找，不会的可以看下202，查找-二分法查找

public int kthSmallest(int[][] matrix, int k) {
    int low = matrix[0][0], high = matrix[matrix.length - 1][matrix[0].length - 1];
    while (low < high) {
        int mid = low + (high - low) / 2;
        int count = 0, j = matrix[0].length - 1;
        for (int i = 0; i < matrix.length; i++) {
            while (j >= 0 && matrix[i][j] > mid)
                j--;
            count += (j + 1);
        }
        if (count < k)
            low = mid + 1;
        else
            high = mid;
    }
    return low;
}

因为矩阵的行和列都是排过序的，这里先找到最中间的值，count表示的是比mid小的值有多少个，每次都是用每行的最右边的一个值和mid比较，直到比mid大，才会执行上面的while循环，然后再往前找，这里的第count个值是大于mid的最小值，所以我们不能在count==k的时候直接return。