LintCode 401. 排序矩阵中的从小到大第k个数
解
第一步，万年不变的查错。如果给的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时，或空间要求常数时，二分法比较好。
