【优先队列】LeetCode 378. 有序矩阵中第 K 小的元素

题目链接

378. 有序矩阵中第 K 小的元素

思路

因为矩阵的每行和每列元素均按升序排序，所以我们可以打破传统思路

代码

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        PriorityQueue<int[]> priorityQueue = new PriorityQueue<>(((o1, o2) -> matrix[o1[0]][o1[1]]
                - matrix[o2[0]][o2[1]]));

        for(int i = 0; i < matrix.length; i++){
            priorityQueue.add(new int[]{i, 0});
        }

        // --k > 0 只循环 k-1 次,k-- > 0 是循环 k 次
        // 循环 k-1 次而不是 k 次的目的是为了在循环结束之后正好弹出第 k 个元素
        while(--k > 0){
            int[] min = priorityQueue.poll();
            int x = min[0];
            int y = min[1] + 1;
            if(y < matrix.length){
                priorityQueue.add(new int[]{x, y});
            }
        }
        int[] result = priorityQueue.poll();

        return matrix[result[0]][result[1]];
    }

}