378. Kth Smallest Element in a Sorted Matrix

Leetcode Binary Search Heap

Given a $n \times n$ matrix where each of the rows and columns are sorted in ascending order, find the $k$th 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 \le k \le n^2$.

分析

使用最小二叉堆，把元素放入堆中，然后依次取出$k$次，第$k$次取出的元素即为第$k$个最小元素。问题是怎么放元素，才能使时间和空间复杂度达到最优。最直接的方法是把所有元素放入最大二叉堆中，最大二叉堆的大小为k，当二叉堆放不下时，取出堆顶元素，当循环结束时，最大二叉堆的堆顶元素恰好是第$k$个最小元素。但这种方法的时间复杂度达到$O(n\log n)$，很明显是不行的：如果把所有元素排序，然后取出第$k$个最小元素，时间复杂度相同。

上面的方法根本没有利用题目的性质：排序的矩阵。也就是说对于矩阵元素$(i, j)$，元素$(i+1, j)$和元素$(i, j+1)$都大于$(i, j)$。利用这样的性质，可以设想，把第一行元素都放入最小二叉堆中，然后每次取出一个元素$(i, j)$，将大于该元素的元素$(i+1, j)$放入到二叉堆中，依次循环$k$次，第$k$次取出的元素即为第$k$个最小元素。

public int kthSmallest(int[][] matrix, int k) {
    PriorityQueue<Tuple> pq = new PriorityQueue<>();
    int n = matrix.length;
    for (int i = 0; i < n; i++)
        pq.offer(new Tuple(matrix[0][i], 0, i));
    Tuple cur;
    for (int i = 1; i < k; i++) {
        cur = pq.poll();
        if (cur.x + 1 < n) pq.offer(new Tuple(matrix[cur.x + 1][cur.y], cur.x + 1, cur.y));     
    }
    return pq.poll().val;
}

class Tuple implements Comparable<Tuple> {
    int x;
    int y;
    int val;
    public Tuple(int val, int x, int y) {
        this.x = x;
        this.y = y;
        this.val = val;
    }

    public int compareTo(Tuple another) {
        return val - another.val;
    }

}