class Solution {
public:
vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
}
};
剑指 Offer II 061. 和最小的 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]
示例 3:
输入: nums1 = [1,2], nums2 = [3], k = 3 输出: [1,3],[2,3] 解释: 也可能序列中所有的数对都被返回:[1,3],[2,3]
提示:
1 <= nums1.length, nums2.length <= 104-109 <= nums1[i], nums2[i] <= 109nums1, nums2 均为升序排列1 <= k <= 1000
注意：本题与主站 373 题相同：https://leetcode.cn/problems/find-k-pairs-with-smallest-sums/
原站题解
golang 解法, 执行用时: 0 ms, 内存消耗: 3.4 MB, 提交时间: 2022-11-22 10:15:24
func kSmallestPairs(nums1, nums2 []int, k int) (ans [][]int) {
m, n := len(nums1), len(nums2)
// 二分查找第 k 小的数对和
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
})
// 找数对和小于 pairSum 的数对
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
}
}
}
// 找数对和等于 pairSum 的数对
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
}
python3 解法, 执行用时: 44 ms, 内存消耗: 15.3 MB, 提交时间: 2022-11-22 10:14:36
class Solution:
def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
m, n = len(nums1), len(nums2)
# 二分查找第 k 小的数对和
left, right = nums1[0] + nums2[0], nums1[m - 1] + nums2[n - 1] + 1
while left < right:
mid = (left + right) // 2
cnt = 0
i, j = 0, n - 1
while i < m and j >= 0:
if nums1[i] + nums2[j] > mid:
j -= 1
else:
cnt += j + 1
i += 1
if cnt < k:
left = mid + 1
else:
right = mid
pairSum = left
ans = []
# 找数对和小于 pairSum 的数对
i = n - 1
for num1 in nums1:
while i >= 0 and num1 + nums2[i] >= pairSum:
i -= 1
for j in range(i + 1):
ans.append([num1, nums2[j]])
if len(ans) == k:
return ans
# 找数对和等于 pairSum 的数对
i = n - 1
for num1 in nums1:
while i >= 0 and num1 + nums2[i] > pairSum:
i -= 1
j = i
while j >= 0 and num1 + nums2[j] == pairSum:
ans.append([num1, nums2[j]])
if len(ans) == k:
return ans
j -= 1
return ans
golang 解法, 执行用时: 8 ms, 内存消耗: 3.7 MB, 提交时间: 2022-11-22 10:07:13
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 }
python3 解法, 执行用时: 44 ms, 内存消耗: 15.3 MB, 提交时间: 2022-11-22 10:06:53
'''
核心思想是把候选数对加入最小堆，实现排序
'''
class Solution:
def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
from heapq import heappop, heappush
m, n = len(nums1), len(nums2)
ans = []
pq = [(nums1[i] + nums2[0], i, 0) for i in range(min(k, m))]
while pq and len(ans) < k:
_, i, j = heappop(pq)
ans.append([nums1[i], nums2[j]])
if j + 1 < n:
heappush(pq, (nums1[i] + nums2[j + 1], i, j + 1))
return ans