[LeetCode] 378. Kth Smallest Element in a Sorted Matrix

Given a n x n matrix where each of the rows and columns are sorted in ascending order, find the kth smallest element in the matrix.

Note that it is the kth smallest element in the sorted order, not the kth distinct element.

Example:

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

Note:
You may assume k is always valid, 1 ≤ k ≤ n^2.

有序矩阵中第K小的元素。

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

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/kth-smallest-element-in-a-sorted-matrix
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

这道题有两种解法，1是priority queue，2是二分法。

首先是priority queue的做法。如果是无脑将matrix中所有的元素都加入一个最小堆中然后找第K小的元素，时间复杂度会非常高，因为pq的时间就是O(nlogn) + 遍历input的时间O(n^2)，所以时间复杂度起码是O(n^2)起跳。因为matrix里面每一列也是有序的，所以这里优化的方法是将每个元素的坐标和对应的val做成一个tuple加入pq，同时，pq中只有matrix.length个元素进行比较。这里跑一个例子说一下吧。pq一开始加入了1，5，9三个元素，弹出1之后，加入的是10；第二次的比较是介于10，5，9之间的，会弹出5，并且加入跟5同一列的11。以此循环往复K - 1次之后，再下一次poll出来的元素就是第K小的元素了。

时间O(nlogn)

空间O(n)

Java实现

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        PriorityQueue<Tuple> pq = new PriorityQueue<>(matrix.length, (a, b) -> (a.val - b.val));
        for (int i = 0; i < matrix.length; i++) {
            pq.offer(new Tuple(0, i, matrix[0][i]));
        }
        for (int i = 0; i < k - 1; i++) {
            Tuple tuple = pq.poll();
            if (tuple.x == matrix.length - 1) {
                continue;
            }
            pq.offer(new Tuple(tuple.x + 1, tuple.y, matrix[tuple.x + 1][tuple.y]));
        }
        return pq.poll().val;
    }
}

class Tuple {
    int x, y, val;

    public Tuple(int x, int y, int val) {
        this.x = x;
        this.y = y;
        this.val = val;
    }
}

[2021年2月更新] 用最小堆实现的暴力解是可以通过的。

时间O(n^2)

空间O(n^2) - 因为存放了所有的元素

Java实现

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        PriorityQueue<Integer> queue = new PriorityQueue<>();
        for (int i = 0; i < matrix.length; i++) {
            for (int j = 0; j < matrix[0].length; j++) {
                queue.offer(matrix[i][j]);
            }
        }
        while (k > 1) {
            queue.poll();
            k--;
        }
        return queue.peek();
    }
}

二分法的思路有一些特别。一开始找的 left 和 right 是 matrix 里面的最小值和最大值（起码这样确定了值的范围），但是在计算 mid 的时候，mid 的值很有可能不在 matrix 中，因为他不是根据坐标得来的，但是一定处于 left 和 right 之间。所以之后 count() 函数是在计算 matrix 中有多少个元素小于/大于 mid，以此决定到底是 left = mid 还是 right = mid。count() 函数还是比较巧妙的，首先记住 matrix 里面每一行和每一列都是有序的。count() 函数是从 matrix 的左下角那个元素开始看，如果左下角那个元素大于 target（nums[mid]），那么说明左下角元素所在的那一行都不满足，所以要 i--；反之如果左下角那个元素小于 target（nums[mid]），那说明左下角那个元素所在的整列元素都小于 target。如果实在记不住，暴力 O(n^2) 遍历数组去数有多少个数字小于 mid 也行。

跑一下例子吧，首先找到左下角元素12，此时 target 是11，12 > 11，此时只能 i--，因为12所在的那一行一定都不满足（12，13，15）；再来会在第二行看，10 < 11，res += i + 1 => res += 0 + 1 => res += 1。i是10的横坐标0，又因为10上方的所有元素都一定小于10，所以比11小的元素个数是被累加到 res 上去的。

注意：这里的left mid right是数值，不是索引位置

时间O(nlogn)，worse case O(n^2)

空间O(1)

Java实现

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 + 1 < right) {
            int mid = left + (right - left) / 2;
            int num = count(matrix, mid);
            if (num >= k) {
                right = mid;
            } else {
                left = mid;
            }
        }
        if (count(matrix, right) <= k - 1) {
            return right;
        }
        return left;
    }

    private int count(int[][] matrix, int target) {
        int n = matrix.length;
        int res = 0;
        // start from bottom left corner
        int i = n - 1, j = 0;
        while (i >= 0 && j < n) {
            if (matrix[i][j] < target) {
                res += i + 1;
                j++;
            } else {
                i--;
            }
        }
        return res;
    }
}

JavaScript实现

/**
 * @param {number[][]} matrix
 * @param {number} k
 * @return {number}
 */
var kthSmallest = function (matrix, k) {
    let lo = matrix[0][0];
    let hi = matrix[matrix.length - 1][matrix[0].length - 1] + 1;
    while (lo < hi) {
        let mid = lo + Math.floor((hi - lo) / 2);
        let count = 0;
        for (let i = 0; i < matrix.length; i++) {
            for (let j = 0; j < matrix.length; j++) {
                if (matrix[i][j] <= mid) {
                    count++;
                } else {
                    break;
                }
            }
        }
        if (count < k) {
            lo = mid + 1;
        } else {
            hi = mid;
        }
    }
    return lo;
};

相关题目

373. Find K Pairs with Smallest Sums

378. Kth Smallest Element in a Sorted Matrix

719. Find K-th Smallest Pair Distance

LeetCode 题目总结