373. Find K Pairs with Smallest Sums
有兩個array nums1、nums2以遞增的順序排序
定義分別從nums1、nums2各取一個元素所組成的組合[nums1[i]、nums2[j]]
請回傳前k總和最小的組合
思路:
因為nums1、nums2是以遞增順序排列
所以最小的組合一定是[nums1[0]、nums2[0]]
那第二小的組合會是[nums1[0]、nums2[1]] or [nums1[1]、nums2[0]]其中一個
以此類推[nums1[i]、nums2[j]]下一個的組合會是
[nums1[i+1]、nums2[j]] or [nums1[i]、nums2[j+1]]中的一個
知道了這個關係，我們就可以用heap來找到k個最小的組合
然後要記得紀錄每個放進去heap的組合，避免重複放入
golang code :
type h [][2]int
var arr1 []int
var arr2 []int
func (this h) Swap(i, j int) {
this[i], this[j] = this[j], this[i]
}
func (this h) Less(i, j int) bool {
return arr1[this[i][0]]+arr2[this[i][1]] < arr1[this[j][0]]+arr2[this[j][1]]
}
func (this h) Len() int {
return len(this)
}
func (this *h) Push(x interface{}) {
(*this) = append(*this, x.([2]int))
}
func (this *h) Pop() interface{} {
n := len(*this) - 1
res := (*this)[n]
(*this) = (*this)[:n]
return res
}
func kSmallestPairs(nums1 []int, nums2 []int, k int) [][]int {
rec := make(map[[2]int]bool)
arr1 = nums1
arr2 = nums2
ans := make([][]int, 0)
hp := h([][2]int{})
heap.Push(&hp, [2]int{0, 0})
for k > 0 {
temp := heap.Pop(&hp).([2]int)
ans = append(ans, []int{nums1[temp[0]], nums2[temp[1]]})
if temp[1]+1 < len(nums2) && !rec[[2]int{temp[0], temp[1] + 1}] {
heap.Push(&hp, [2]int{temp[0], temp[1] + 1})
rec[[2]int{temp[0], temp[1] + 1}] = true
}
if temp[0]+1 < len(nums1) && !rec[[2]int{temp[0] + 1, temp[1]}] {
heap.Push(&hp, [2]int{temp[0] + 1, temp[1]})
rec[[2]int{temp[0] + 1, temp[1]}] = true
}
k