1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
| class Solution {
private static class HeapNode {
int row;
int col;
int num;
HeapNode(int _row, int _col, int _num) {
row = _row;
col = _col;
num = _num;
}
}
public int kthSmallest(int[][] matrix, int k) {
int initialHeapSize = Math.min(matrix.length, k);
PriorityQueue<HeapNode> minHeap = new PriorityQueue<>((a, b) -> a.num - b.num);
for (int i = 0; i < initialHeapSize; i++) {
minHeap.offer(new HeapNode(i, 0, matrix[i][0]));
}
int count = 0;
while (count < k - 1) {
HeapNode currNode = minHeap.poll();
int currRow = currNode.row, currCol = currNode.col;
if (currCol + 1 < matrix[0].length) {
minHeap.offer(new HeapNode(currRow, currCol + 1, matrix[currRow][currCol + 1]));
}
count += 1;
}
return minHeap.peek().num;
}
}
|
N sorted lists求第K小. 这种类型题的解题思路要知道就是用Heap. 我们用heap存N个行每行的第0个元素, 然后开始poll. poll出来的我们查它所在的行和列, 把这一行下一个元素再请到heap里. 以此类推. 当我们poll了k - 1次时, 此时heap的顶就是答案. 由于我们是每行每列都sort了, 那么我们有时候没必要把每一行的第0个元素都装进heap里面, 比如如果有100行, k只是5, 那么前五行的第0个元素装进heap就行了. 从第6行到最后一行的所有元素都不可能是答案. 因为第六行到最后一行里面最小的就是第六行第0个元素. 而第六行第0个元素又比第五行第0个元素大. 因此前五行的第0个元素其一当答案都不会第六行第0个元素当答案.
时间复杂度: 装heap那就是log(1) + log2 + log3 +…log(min(n, k)) = min(n, k)log(min(n, k)) log(n!)收敛于nlogn. 就是min(n, k)是m, 那么装heap这一步就是O(mlogm). 然后每次poll再装是O(1) + O(logm), 一共K - 1次, 那就是O(Klogm), 加在一起就是O((K + m)logm).
空间复杂度: O(m) 也就是heap的大小.
