【Lintcode】401. Kth Smallest Number in Sorted Matrix

题目地址：

https://www.lintcode.com/problem/kth-smallest-number-in-sorted-matrix/description

给定一个二维矩阵，每行从左到右是递增的，每列从上到下也是递增的。求该矩阵的第$k$小数。

这道题可以用有向图的最短路径模型来思考。首先很显然最左上角的数就是最小的数，我们假想从点$(x,y)$到点$(x+1,y)$和$(x,y+1)$都有一条边，边权是后者与前者上的值的差，很显然边权非负。于是，问题可以转化为，求以$(0,0)$为源的单源最短路问题，求的是离源点的第$k$近的点（第$1$近的点为源点自己）。所以，这就可以用Dijkstra算法来解决。

但需要注意的是，这里每个点只需进堆一次，因为其与源点的距离并不会被更新第二次，也就是说，某个点第一次遍历到的时候，与源点的最短路就已经被确定了，后面是不会被再次更新的。所以当一个顶点入堆的时候，就可以立即标记为visited，下次不必再入一次。代码如下：

import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;

public class Solution {
    class Point {
        int x, y, value;
        Point(int x, int y, int value) {
            this.x = x;
            this.y = y;
            this.value = value;
        }
    
        @Override
        public boolean equals(Object o) {
            Point point = (Point) o;
            return x == point.x && y == point.y;
        }
    
        @Override
        public int hashCode() {
            return 31 * x + y;
        }
    }
    /**
     * @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) {
        // write your code here
        PriorityQueue<Point> minHeap = new PriorityQueue<>((p1, p2) -> p1.value <= p2.value ? -1 : 1);
        Set<Point> visited = new HashSet<>();
        Point p0 = new Point(0, 0, matrix[0][0]);
        minHeap.offer(p0);
        visited.add(p0);
        // count记录出堆的数是第几小的数
        int count = 0;
        while (!minHeap.isEmpty()) {
            Point cur = minHeap.poll();
            count++;
            // 如果出堆的数恰好是第k小的数，则直接将其返回
            if (count == k) {
                return cur.value;
            }
            // 将cur的未访问过的且下一步能到达的邻居节点得到并加入堆中
            for (Point point : generate(cur, matrix, visited)) {
                minHeap.offer(point);
            }
        }
        
        // 这一步不会执行到。如果执行到，则说明k为负或者k超出了矩阵的规模
        return Integer.MAX_VALUE;
    }
    
    private Set<Point> generate(Point z, int[][] matrix, Set<Point> visited) {
        Set<Point> neighbors = new HashSet<>();
        int[][] dirs = {{0, 1}, {1, 0}};
        for (int i = 0; i < dirs.length; i++) {
            int x = z.x + dirs[i][0];
            int y = z.y + dirs[i][1];
            if (0 <= x && x < matrix.length && 0 <= y && y < matrix[0].length) {
                Point neighbor = new Point(x, y, matrix[x][y]);
                if (!visited.contains(neighbor)) {
                    neighbors.add(neighbor);
                    visited.add(neighbor);
                }
            }
        }
        
        return neighbors;
    }
}

时间复杂度$O(k\log k)$，空间$O(k)$。