虽然这里的两个数组都是升序排列,但是两个数组的数对和的大小是不确定的,所以用暴力扫描是无法实现的。最小或者最大的K个值一般都用最小堆或者最大堆。
这里可以用最小堆。需要注意的:
- 堆内不存数对和的大小,而是存
nums1和nums2的索引 compare比较的是nums1和nums2数对和的大小
class MinBinaryHeap {
constructor(compare) {
this.s = []
this.len = 0
this.compare = compare
}
push(el) {
this.s[this.len] = el
++this.len
if (this.len > 1) {
this.swim(this.len - 1)
}
}
pop() {
if (this.len === 0) {
return null
}
const top = this.s[0]
if (this.len > 1) {
this.s[0] = this.s[this.len - 1]
this.sink(0)
}
this.len--
return top
}
peek() {
if (this.len === 0) {
return null
} else {
return this.s[0]
}
}
sink(idx) {
const lastIdx = this.len - 1
while (idx <= lastIdx) {
const rightIdx = (idx + 1) * 2
const leftIdx = rightIdx - 1
if (leftIdx <= lastIdx) {
let smallValueIdx = leftIdx
if (rightIdx <= lastIdx && this.compare(this.s[smallValueIdx], this.s[rightIdx]) > 0) {
smallValueIdx = rightIdx
}
if (this.compare(this.s[idx], this.s[smallValueIdx]) > 0) {
this.swap(idx, smallValueIdx)
idx = smallValueIdx
} else {
break
}
} else {
break
}
}
}
swim(idx) {
while (idx >= 0) {
let parent = (idx - 1) >> 1
if (parent >= 0 && this.compare(this.s[parent], this.s[idx]) > 0) {
this.swap(parent, idx)
idx = parent
} else {
break
}
}
}
swap(a, b) {
let temp = this.s[a]
this.s[a] = this.s[b]
this.s[b] = temp
}
compare(a, b) {
return a - b
}
}
const kSmallestPairs = function(nums1, nums2, k) {
const newPairsHeap = new MinBinaryHeap((a, b) => nums1[a[0]] + nums2[a[1]] - nums1[b[0]] - nums2[b[1]])
const n1l = nums1.length
const n2l = nums2.length
// 保存k个值
// 可以确定的是 最小的是[0, 0]
for(let i = 0; i < Math.min(k, n1l); i++) {
newPairsHeap.push([i, 0])
}
const ans = []
while(newPairsHeap.len) {
const top = newPairsHeap.pop()
ans.push([nums1[top[0]], nums2[top[1]]])
if (ans.length === k) {
return ans
}
// 索引[0, 0]对应元素被加入到答案之后,将[0, 1]加入堆中,
// [1, 0]之前已经加入了。因为是排过序的,所以最小值肯定是[0, 1]或者[1, 0]。
// 这样索引较小的数对已加入堆中,保证下次为最小值
if ((top[1] + 1) < n2l) {
newPairsHeap.push([top[0], top[1] + 1])
}
}
return ans
};