LintCode第401题-排序矩阵中的从小到大第k个数

描述

在一个排序矩阵中找从小到大的第 k 个整数。

排序矩阵的定义为：每一行递增，每一列也递增。

样例 1:

输入:
[[1 ,5 ,7],[3 ,7 ,8],[4 ,8 ,9],
]
k = 4
输出: 5

样例 2:

输入: 
[[1, 2],[3, 4]
]
k = 3
输出: 3

挑战

时间复杂度 O(klogn), n 是矩阵的宽度和高度的最大值

思路:小顶堆（k 路归并法）

易错点:

需要每次都从当前候选里挑出最小的那个再扩展，而不是先随便收集一些数再统一排序

代码如下:

import java.util.*;

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) {

        if (matrix == null || matrix.length < 1 || matrix[0].length < 1) {

            return 0;

        }

        int n = matrix.length;

        int m = matrix[0].length;

        int smallestNumber = 0; // 最后返回的结果

        int calculateNum = 1;   // 计数器：表示当前弹出的元素个数

        boolean visited[][] = new boolean[n][m];

        // 先进入一个节点 即0,0位置节点

        visited[0][0] = true;

        // 小顶堆：每次取出全局最小的元素

        PriorityQueue<int[]> integerQueue = new PriorityQueue<>(

            (a, b) ->Integer.compare(matrix[a[0]][a[1]],matrix[b[0]][b[1]])

        );

        // 把 (0,0) 放进堆

        integerQueue.offer(new int[]{0, 0});

        // 依次弹出k次，第k次弹出的就是答案

        for (int t = 0; t < k; t++) {

            int[] cur = integerQueue.poll();

            int i = cur[0];

            int j = cur[1];

            smallestNumber = matrix[i][j];

            // 遍历其右边的第一个元素和下面的第一个元素

            if (i + 1 < n && visited[i + 1][j] == false) {

                integerQueue.offer(new int[]{i + 1, j, matrix[i + 1][j]});

                visited[i + 1][j] = true;

                calculateNum++;

            }

            if (j + 1 < m && visited[i][j + 1] == false) {

                integerQueue.offer(new int[]{i, j + 1, matrix[i][j + 1]});

                visited[i][j + 1] = true;

                calculateNum++;

            }

        }

        return smallestNumber;

    }

}