[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<Node> queue = new PriorityQueue<Node>((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;
    }
}