题目地址:
https://www.lintcode.com/problem/kth-smallest-number-in-sorted-matrix/description
给定一个二维矩阵,每行从左到右是递增的,每列从上到下也是递增的。求该矩阵的第 k k k小数。
这道题可以用有向图的最短路径模型来思考。首先很显然最左上角的数就是最小的数,我们假想从点 ( x , y ) (x,y) (x,y)到点 ( x + 1 , y ) (x+1,y) (x+1,y)和 ( x , y + 1 ) (x,y+1) (x,y+1)都有一条边,边权是后者与前者上的值的差,很显然边权非负。于是,问题可以转化为,求以 ( 0 , 0 ) (0,0) (0,0)为源的单源最短路问题,求的是离源点的第 k k k近的点(第 1 1 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\log k) O(klogk),空间 O ( k ) O(k) O(k)。