【leetcode】有序矩阵中第 K 小的元素

一、题目描述

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

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

示例 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

二、代码思路

第一眼来看最简单的解法就是把二维数组的元素都放到一维数组中然后使用快排升序即可。

但是，题目要求空间复杂度低于O(N2)，所以我们可以考虑维护大小为K的优先级队列（大根堆），依次将矩阵元素放入大根堆，一旦大根堆大小超过K那么我们剔除掉当前最大的元素，当遍历完n个元素之后大根堆中就存放着最小的K个元素，而且根节点是第K小元素。n2logk

当然，还有第三种解法，就是归并排序，题目说了矩阵每一行都是有序的，那么我们可以理解为从n个有序的数组中选出第K个有序的数，那么自然而然可以想到归并，但是这里的归并并非是两两归并而是n个数组一起归并，这个也很好归并，从n个数中选最小的即可（借助大小为k的小根堆排序），然后取出来，再放入新的数归并即可。看代码更好理解。

三、代码题解

public class Solution02 {
    public int kthSmallest(int[][] matrix, int k) {
        //优先级队列 + lambda表达式（匿名内部类）小根堆
        PriorityQueue<int[]> queue = new PriorityQueue<>((a,b) -> {
            return a[0] - b[0];
        });
        //向小根堆中加入数据
        for (int i = 0; i < matrix.length; i++) {
            //需要归并的数组的首元素 | 第几个数组 | 这个数组归并到第几个元素了
            queue.offer(new int[] {matrix[i][0], i, 0});
        }
        //开始归并, 这里之所以是k - 1，因为我们要找第k个小的元素，我们先找到前k - 1个，然后最后第k个元素就存放在堆头
        int length = matrix[0].length;
        for (int i = 0; i < k - 1; i++) {
            //找到当前最小的元素
            int[] min = queue.poll();
            int curIndex = min[1];
            int curNum = min[2];
            //需要归并的数组的首元素 | 第几个数组 | 这个数组归并到第几个元素了
            if (curNum < length - 1) {
                queue.offer(new int[] {matrix[curIndex][curNum + 1], curIndex, curNum + 1});
            }
        }
        return queue.poll()[0];
    }
}

时间复杂度： o(klogn) logn是调整根堆 k是调整k次