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

Problem

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.

Solution

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        Queue queue = new PriorityQueue((a, b)->a.val-b.val);
        //先放一行，或一列
        for (int i = 0; i < n; i++) queue.offer(new Node(i, 0, matrix[i][0]));
        //把堆顶的最小元素取出来，取k-1次，如果该node有下一行/下一列的node，放入堆中
        for (int i = 0; i < k-1; i++) {
            Node node = queue.poll();
            if (node.y == n-1) continue;
            else queue.offer(new Node(node.x, node.y+1, matrix[node.x][node.y+1]));
        }
        //最小的k-1个元素已经在上面的for循环被poll完了，下一个堆顶元素就是kth smallest
        return queue.poll().val;
    }
}

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