100.查找和最小的 K 对数字
题目
给定两个以 非递减顺序排列 的整数数组 nums1 和 nums2 , 以及一个整数 k 。
定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2 。
请找到和最小的 k 个数对 (u1,v1), (u2,v2) ... (uk,vk) 。
示例 1:
- 输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
- 输出: [[1,2],[1,4],[1,6]] 解释: 返回序列中的前 3 对数： [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
示例 2:
- 输入: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
- 输出: [[1,1],[1,1]] 解释: 返回序列中的前 2 对数： [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]
思路
最直接的暴力方法
假设:
nums1 = [1, 7, 11]
nums2 = [2, 4, 6]
我们可以把所有组合都枚举出来：
(1,2) → 3
(1,4) → 5
(1,6) → 7
(7,2) → 9
(7,4) → 11
(7,6) → 13
(11,2) → 13
(11,4) → 15
(11,6) → 17
然后排序，取前 K 个。
如果两个数组长度分别为 m 和 n：
- 一共有 m * n 个组合
- 排序需要 O(mn log(mn))
显然不够优秀。
我们真正需要思考的是：
两个数组本身已经有序，能不能利用这个性质？
答案是可以。
假设:
nums1 = [1, 7, 11]
nums2 = [2, 4, 6]
我们可以把它们看成三条有序的「链」：
以 1 开头：
[1,2] → 3 [1,4] → 5 [1,6] → 7
以 7 开头：
[7,2] → 9 [7,4] → 11 [7,6] → 13
以11开头:
[11,2] → 13 [11,4] → 15 [11,6] → 17
注意，因为两个数组都是从小到大排列的，所以每一行都是递增的。
现在问题就变成了：
从多个已经排好序的序列中，找出最小的K个元素。
对于 nums1 中的每一个元素 a，它和 nums2 中的所有元素配对，这些配对的和是随着 nums2 下标增加而递增的（因为 nums2 本身升序）。
这其实就是非常经典的:
多路归并(K-way Merge)
而多路归并最经典的数据结构就是:
最小堆
那我们需要不需要把上面例子中所有的9个数字都计算出来？ 其实： 不需要，因为每一排本身已经是从小到大排列的。
比如第一排:
3 -> 5 -> 7，所以一开始我们只需要关注3就可以，因为5和7一定都排在3后面
第二排:
9 -> 11 -> 13，我们只需要先关注9。
第三排：
13 -> 15 -> 17，我们只需要先关注13.
所以一开始，我们只需要三个数字(每一排的第一个数):
3
9
13
我们现在手里有:
3
9
13
谁最小？
3最小，所以第一个答案就是[1, 2]，
我们把3拿走之后怎么办？
把第一排往后移动一步，第一排原来是3 -> 5 -> 7，我们拿走了3，那么第一排现在应该看: 5。
于是现在我们关注:
5
9
13
谁最小？ 5，所以第二个答案就是[1, 4]，再把第一排向后移动。
整个过程
第一排：3 → 5 → 7
第二排：9 → 11 → 13
第三排：13 → 15 → 17
我们开始每次只看每一排的第一个数:
3
9
13
然后找这几个当前数字里面最小的那个。
找到了之后： 把这一排往后移动一个，然后继续找最小的。
但是假设有10000个数组，我们每次都需要再10000个数字里面找最小的。
如果找很多次，又会变慢，所以我们需要一个东西:
自动帮我们维护当前最小值。
这就是PriorityQueue，也就是小顶堆。
那么PriorityQueue里面到底放什么？
我们不能只放:
3
9
13
因为我们不仅需要知道: 当前最小的和是多少。
还需要知道： 这个最小的和来自哪两个数字？
例如： 3其实是来自nums1[0]和nums2[0]的和。
所以我们在PriorityQueue里面存的是: [i, j]，也就是两个数组的下标。
也就是说: [i, j]表示的是nums1[i] 和 nums2[j]。
这就是我们要写PriorityQueue的Comparator的地方:
val pq = PriorityQueue<IntArray> { a, b ->
(nums1[a[0]] + nums2[a[1]])
.compareTo(nums1[b[0]] + nums2[b[1]])
}
也就是告诉PriorityQueue：
我这里面放的是int数组[i, j]，你不要直接比较[i, j]，你要计算nums1[i] = nums2[j]，谁的和小，谁就优先。
代码实现
class Solution {
fun kSmallestPairs(
nums1: IntArray,
nums2: IntArray,
k: Int
): List<List<Int>> {
val result = mutableListOf<List<Int>>()
if (nums1.isEmpty() || nums2.isEmpty()) {
return result
}
// 小顶堆
val pq = PriorityQueue<IntArray> { a, b ->
(nums1[a[0]] + nums2[a[1]])
.compareTo(nums1[b[0]] + nums2[b[1]])
}
// 每一排先放第一个
val count = minOf(nums1.size, k)
for (i in 0 until count) {
pq.offer(intArrayOf(i, 0))
}
// 找 k 个
repeat(k) {
if (pq.isEmpty()) {
return@repeat
}
// 拿当前最小的
val cur = pq.poll()
val i = cur[0]
val j = cur[1]
// 加入答案
result.add(
listOf(nums1[i], nums2[j])
)
// 这一排还有下一个
if (j + 1 < nums2.size) {
pq.offer(
intArrayOf(i, j + 1)
)
}
}
return result
}
}
PriorityQueue介绍
PriorityQueue是Java/Kotlin一个非常重要的数据结构。
一个自动帮你维护最值的容器。
最常见的用途就是:
- 找最小值 -> 小顶堆
- 找最大值 -> 大顶堆
- Top K
- K路合并
- 贪心算法
- BFS等
PriorityQueue其实就是一个队列。
普通队列是: 先进先出 FIFO 但是PriorityQueue不一样，它是:
谁的优先级最高，谁先出来。
例如:
val pq = PriorityQueue<Int>()
pq.offer(5)
pq.offer(2)
pq.offer(8)
pq.offer(1)
虽然加入的顺序是: 5, 2, 8, 1 但是: pq.poll()，会得到1，再poll()是2，再poll是5，最后是8.
所以PriorityQueue会按照优先级取元素，默认是最小的元素优先。
最重要的三个方法:
- offer(): 加入元素
- peek(): 查看堆顶，但是不删除
- poll(): 取出对顶，同时删除它。
那最大堆怎么办？
因为PriorityQueue默认是小顶堆，如果想要大顶堆，那就需要告诉PriorityQueue:
数字越大，优先级越高
最简单的写法是:
val pq = PriorityQueue<Int>(compareByDescending { it })
另一种写法就是: Comparator: 自定义优先级规则
val pq = PriorityQueue<Int> { a, b ->
b.compareTo(a)
}
复杂度
- 时间复杂度: O(k log(min(m, n)))
- 空间复杂度: O(min(m, n))（堆的大小）
- 邮箱 ：charon.chui@gmail.com
- Good Luck!