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