剑指 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] <= 109
nums1, nums2 均为升序排列
1 <= k <= 1000

维护一个大小为k的大顶堆，存储数对及其和

class Solution:
    def kSmallestPairs(self, nums1: list[int], nums2: list[int], k: int) -> list[list[int]]:
        a, b = len(nums1), len(nums2)
        k = min(k, a * b)
        a, b = min(a, k), min(b, k)
        minHeap = []
        for i in range(a):
            for j in range(b):
                tmp = nums1[i] + nums2[j]
                heapq.heappush(minHeap, (-tmp, nums1[i], nums2[j]))
                while len(minHeap) > k:
                    heapq.heappop(minHeap)
        ans = []
        while minHeap:
            res = heapq.heappop(minHeap)
            ans.append([res[1], res[2]])
        return ans

type hp [][3]int

func (h hp) Len() int {
	return len(h)
}

func (h hp) Less(i, j int) bool {
	return h[i][2] > h[j][2]
}

func (h *hp) Push(i interface{}) {
	*h = append(*h, i.([3]int))
}

func (h *hp) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[:n-1]
	return x
}

func (h hp) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func kSmallestPairs(nums1 []int, nums2 []int, k int) [][]int {
	a, b := len(nums1), len(nums2)
	k = min(k, a*b)
	a, b = min(a, k), min(b, k)
	minHeap := &hp{}
	heap.Init(minHeap)
	for i := 0; i < a; i++ {
		for j := 0; j < b; j++ {
			heap.Push(minHeap, [3]int{nums1[i], nums2[j], nums1[i] + nums2[j]})
			for minHeap.Len() > k {
				heap.Pop(minHeap)
			}
		}
	}
	ret := [][]int{}
	for i := 0; i < k; i++ {
		res := heap.Pop(minHeap).([3]int)
		a := []int{res[0], res[1]}
		ret = append(ret, a)
	}
	return ret
}