public int kthSmallest(int[][] matrix, int k) {
if (matrix == null || matrix.length == 0 || matrix[0].length == 0 || k < 0) {
return Integer.MAX_VALUE;
}
Queue<Loc> pq = new PriorityQueue<>();
HashSet<Loc> visited = new HashSet<>();
LinkedList<Integer> res = new LinkedList<>();
Loc first = new Loc(0, 0, matrix[0][0]);
pq.offer(first);
visited.add(first);
while (res.size() < k) {
Loc tmp = pq.poll();
res.add(tmp.val);
int nextI = tmp.i + 1;
int nextJ = tmp.j + 1;
if (nextI < matrix.length) {
Loc next1 = new Loc(nextI, tmp.j, matrix[nextI][tmp.j]);
if (!visited.contains(next1)) {
pq.offer(next1);
visited.add(next1);
}
}
if (nextJ < matrix[tmp.i].length) {
Loc next2 = new Loc(tmp.i, nextJ, matrix[tmp.i][nextJ]);
if (!visited.contains(next2)) {
pq.offer(next2);
visited.add(next2);
}
}
}
return res.getLast();
}
class Loc implements Comparable<Loc> {
int i;
int j;
int val;
public Loc(int iLoc, int jLoc, int v) {
i = iLoc;
j = jLoc;
val = v;
}
public int compareTo(Loc o) {
return val - o.val;
}
@Override
public String toString() {
return String.valueOf(val);
}
@Override
public boolean equals(Object obj) {
if (obj instanceof Loc) {
Loc tmp = (Loc) obj;
return i == tmp.i && j == tmp.j && val == tmp.val;
}
return false;
}
@Override
public int hashCode() {
return Objects.hash(i, j, val);
}
}
简洁d版:
class Solution {
public int kthSmallest(int[][] matrix, int k) {
if (matrix == null || matrix.length == 0 || matrix[0].length == 0 || k < 0) {
return -1;
}
int n = matrix.length;
int m = matrix[0].length;
PriorityQueue<Loc> minHeap = new PriorityQueue<>(k, new Comparator<Loc>(){
public int compare(Loc l1, Loc l2){
return l1.val - l2.val;
}
});
boolean[][] visited = new boolean[n][m];
LinkedList<Integer> res = new LinkedList<>();
minHeap.offer(new Loc(0, 0, matrix[0][0]));
visited[0][0] = true;
while (res.size() < k) {
Loc cur = minHeap.poll();
res.add(cur.val);
int nexti = cur.i + 1;
int nextj = cur.j + 1;
if (nextj < m && !visited[cur.i][nextj]) {
minHeap.offer(new Loc(cur.i, nextj, matrix[cur.i][nextj]));
visited[cur.i][nextj] = true;
}
if (nexti < n && !visited[nexti][cur.j]) {
minHeap.offer(new Loc(nexti, cur.j, matrix[nexti][cur.j]));
visited[nexti][cur.j] = true;
}
}
return res.getLast();
}
}
class Loc {
int i;
int j;
int val;
public Loc(int i, int j, int v) {
this.i = i;
this.j = j;
this.val = v;
}
}