
用堆实现优先队列
- 我们可以先将 nums1 的前 k 个索引数对加入序列,然后每次去除元素时,我们只需要将 nums2 的索引增加即可,这样就可以很好的避免重复加入元素的问题
func kSmallestPairs(nums1, nums2 []int, k int) (ans [][]int) {
m, n := len(nums1), len(nums2)
h := hp{nil, nums1, nums2}
for i := 0; i < k && i < m; i++ {
h.data = append(h.data, pair{i, 0})
}
for h.Len() > 0 && len(ans) < k {
p := heap.Pop(&h).(pair)
i, j := p.i, p.j
ans = append(ans, []int{nums1[i], nums2[j]})
if j+1 < n {
heap.Push(&h, pair{i, j + 1})
}
}
return
}
type pair struct{ i, j int }
type hp struct {
data []pair
nums1, nums2 []int
}
func (h hp) Len() int { return len(h.data) }
func (h hp) Less(i, j int) bool { a, b := h.data[i], h.data[j]; return h.nums1[a.i]+h.nums2[a.j] < h.nums1[b.i]+h.nums2[b.j] }
func (h hp) Swap(i, j int) { h.data[i], h.data[j] = h.data[j], h.data[i] }
func (h *hp) Push(v interface{}) { h.data = append(h.data, v.(pair)) }
func (h *hp) Pop() interface{} { a := h.data; v := a[len(a)-1]; h.data = a[:len(a)-1]; return v }
二分查找:
func kSmallestPairs(nums1, nums2 []int, k int) (ans [][]int) {
m, n := len(nums1), len(nums2)
left, right := nums1[0]+nums2[0], nums1[m-1]+nums2[n-1]+1
pairSum := left + sort.Search(right-left, func(sum int) bool {
sum += left
cnt := 0
i, j := 0, n-1
for i < m && j >= 0 {
if nums1[i]+nums2[j] > sum {
j--
} else {
cnt += j + 1
i++
}
}
return cnt >= k
})
i := n - 1
for _, num1 := range nums1 {
for i >= 0 && num1+nums2[i] >= pairSum {
i--
}
for _, num2 := range nums2[:i+1] {
ans = append(ans, []int{num1, num2})
if len(ans) == k {
return
}
}
}
i = n - 1
for _, num1 := range nums1 {
for i >= 0 && num1+nums2[i] > pairSum {
i--
}
for j := i; j >= 0 && num1+nums2[j] == pairSum; j-- {
ans = append(ans, []int{num1, nums2[j]})
if len(ans) == k {
return
}
}
}
return
}