401. 排序矩阵中的从小到大第k个数

描述

在一个排序矩阵中找从小到大的第k个整数。
排序矩阵的定义为：每一行递增，每一列也递增。

样例

给出k=4和一个排序矩阵：

[
  [1 ,5 ,7],
  [3 ,7 ,8],
  [4 ,8 ,9],
]

返回5。

挑战

使用O(k log n)的方法，n为矩阵的宽度和高度中的最大值。

思路

矩阵为行列从小到大排序，左上角即为整个矩阵最小元素，加入最小堆，堆顶元素抛出，每抛出一个元素把当前元素的右边元素和下边元素加入堆，求从小到大排列第K小元素，最小堆抛出 k - 1 次后堆顶即为第 k 小元素

代码

时间复杂度KlogK

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

// 最小堆
class PairComparator implements Comparator {
    public int compare(Pair a, Pair b) {
        return a.val - b.val;
    }
}

public class Solution {
    /**
     * @param matrix: a matrix of integers
     * @param k: an integer
     * @return: the kth smallest number in the matrix
     */
    public int kthSmallest(int[][] matrix, int k) {
        int[] dx = new int[]{0, 1};
        int[] dy = new int[]{1, 0};
        int n = matrix.length;
        int m = matrix[0].length;
        // 布尔类型，默认为 false
        boolean[][] hash = new boolean[n][m];
        PriorityQueue minHeap = new PriorityQueue(k, new PairComparator());
        minHeap.add(new Pair(0, 0, matrix[0][0]));

        // 抛出 k-1 次，第 k 次在栈顶
        for (int i = 0; i < k - 1; i ++) {
            Pair cur = minHeap.poll();
            for (int j = 0; j < 2; j ++){
                int next_x = cur.x + dx[j];
                int next_y = cur.y + dy[j];
                Pair next_Pair = new Pair(next_x, next_y, 0);
                // 没有越界且未曾被添加到PriorityQueue过
                // 堆每抛出一个，将当前抛出元素的下边元素和右边元素加入PriorityQueue
                if (next_x < n && next_y < m && !hash[next_x][next_y]) {
                    hash[next_x][next_y] = true;
                    next_Pair.val = matrix[next_x][next_y];
                    minHeap.add(next_Pair);
                }
            }
        }
        return minHeap.peek().val;
    }
}