46 lines
858 B
Go
46 lines
858 B
Go
package q373
|
|
|
|
import "container/heap"
|
|
|
|
type MHeap [][]int
|
|
|
|
func (m *MHeap) Len() int { return len(*m) }
|
|
|
|
func (m *MHeap) Less(i int, j int) bool {
|
|
return (*m)[i][0]+(*m)[i][1] > (*m)[j][0]+(*m)[j][1]
|
|
}
|
|
|
|
func (m *MHeap) Pop() any {
|
|
ret := (*m)[len(*m)-1]
|
|
*m = (*m)[:len(*m)-1]
|
|
return ret
|
|
}
|
|
|
|
func (m *MHeap) Push(x any) {
|
|
elem, _ := x.([]int)
|
|
*m = append(*m, elem)
|
|
}
|
|
|
|
func (m *MHeap) Swap(i int, j int) { (*m)[i], (*m)[j] = (*m)[j], (*m)[i] }
|
|
|
|
func kSmallestPairs(nums1 []int, nums2 []int, k int) [][]int {
|
|
h := make(MHeap, 0, k)
|
|
for i := range min(k, len(nums1)) {
|
|
for j := range min(k, len(nums2)) {
|
|
if h.Len() < k {
|
|
heap.Push(&h, []int{nums1[i], nums2[j]})
|
|
continue
|
|
}
|
|
|
|
if nums1[i]+nums2[j] >= h[0][0]+h[0][1] {
|
|
break
|
|
}
|
|
h[0][0], h[0][1] = nums1[i], nums2[j]
|
|
heap.Fix(&h, 0)
|
|
}
|
|
}
|
|
|
|
return h
|
|
}
|
|
|
|
var _ = kSmallestPairs
|