解
第一步,万年不变的查错。如果给的array是null或空,那么直接return。
public int kthSmallest(int[][] matrix, int k) {
if (matrix == null || matrix.length == 0 || matrix[0].length == 0) {
return 0;
}
...
}
第一个想法是最暴力的,用MaxHeap把k个数放进去,对比,比它大的不要。到最后MaxHeap里面最大的就是第K小的。
直接上code
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 == 0 || matrix[0].length == 0) {
return 0;
}
Queue<Integer> pq = new PriorityQueue<>((a, b) -> -Integer.compare(a, b));
for (int i = 0; i < matrix.length; i++) {
for (int j = 0; j < matrix[0].length; j++) {
if (pq.size() < k) {
pq.offer(matrix[i][j]);
} else {
if (matrix[i][j] < pq.peek()) {
pq.poll();
pq.offer(matrix[i][j]);
}
}
}
}
return pq.poll();
}
}
但是这个方法并没有用到这个matrix是横竖sorted的特性。肯定不是最优解。
解2
还是用heap,不过这次用MinHeap。然后从左上角开始。左上角的数字,肯定是最小的,因为这是个排序矩阵。然后下一个最小的,肯定在它下面或是右面。把这2个数字都扔进heap,在poll出最小的那个,继续把它下面和右面的数字扔进去。大致的思路就是这样。先建一个Pair用来存放坐标。
private class Pair {
int x;
int y;
public Pair(int x, int y) {
this.x = x;
this.y = y;
}
}
这里九章存了val,但是我认为val可以从matrix里面拿到。
然后初始化一下又来标记下一个位置的数组。
int[] dx = new int[] {0, 1};
int[] dy = new int[] {1, 0};
int n = matrix.length;
int m = matrix[0].length;
还需要建一个hash,这样可以保证已经放进heap的不会再放一次。
boolean[][] hash = new boolean[n][m];
然后开始建一个PriorityQueue,这次用MinHeap,poll n次,就是k smallest的数。
Queue<Pair> pq = new PriorityQueue<>(Comparator.comparingInt(a -> matrix[a.x][a.y]));
从左上角开始,因为它肯定是最小的数。
pq.add(new Pair(0, 0));
然后for循环,poll出当前最小的数,然后把下两个放进heap。
for (int i = 0; i < k - 1; i++) {
Pair current = pq.poll();
for (int j = 0; j < 2; j++) {
int nextX = current.x + dx[j];
int nextY = current.y + dy[j];
Pair nextPair = new Pair(nextX, nextY);
if (nextX < n && nextY < m && !hash[nextX][nextY]) {
hash[nextX][nextY] = true;
pq.add(nextPair);
}
}
}
这里重要是的,for循环k-1次,这样就poll出k-1个最小的数字,然后剩下heap里第一个就是第k小的数。
return matrix[pq.peek().x][pq.peek().y];
完整的code
public class Solution {
/*
* @param matrix: a matrix of integers
* @param k: An integer
* @return: the kth smallest number in the matrix
*/
private class Pair {
int x;
int y;
public Pair(int x, int y) {
this.x = x;
this.y = y;
}
}
public int kthSmallest(int[][] matrix, int k) {
if (matrix == null || matrix.length == 0 || matrix[0].length == 0) {
return 0;
}
int[] dx = new int[] {0, 1};
int[] dy = new int[] {1, 0};
int n = matrix.length;
int m = matrix[0].length;
boolean[][] hash = new boolean[n][m];
Queue<Pair> pq = new PriorityQueue<>(Comparator.comparingInt(a -> matrix[a.x][a.y]));
pq.add(new Pair(0, 0));
for (int i = 0; i < k - 1; i++) {
Pair current = pq.poll();
for (int j = 0; j < 2; j++) {
int nextX = current.x + dx[j];
int nextY = current.y + dy[j];
Pair nextPair = new Pair(nextX, nextY);
if (nextX < n && nextY < m && !hash[nextX][nextY]) {
hash[nextX][nextY] = true;
pq.add(nextPair);
}
}
}
return matrix[pq.peek().x][pq.peek().y];
}
}
分析
第一个解法,heap是size k,读n个数,所以时间是O(nlogk)。空间是O(k)。
第二个解法,循环k遍,poll k个数,每次往里放最多2个数,所以heap最大是2k,hash最大也是2k,即空间O(4k),也就是O(k)。poll了k次,放了最多2k次,每次最多logk,总共时间就是O(klogk)。
总结一下,第一个解法适合一开始给出,然后优化成第二种解法。九章还有一种二分法的解法,在最大和最小之间进行二分,然后判断中间数在数组里排第几。这样,空间是O(1),二分是log(max - min),判断中间数是第几个要O(n)。那总共就是O(nlog(max-min))。当max-min小于k时,或空间要求常数时,二分法比较好。
网友评论