leetcode-go/leetcode/medium/373.find-k-pairs-with-smallest-sums.go

type pair struct {
	_1 int
	_2 int
}

type pairs []pair

func (ps pairs) Len() int           { return len(ps) }
func (ps pairs) Less(i, j int) bool { return ps[i]._1 < ps[j]._1 }
func (ps pairs) Swap(i, j int)      { ps[i], ps[j] = ps[j], ps[i] }

func (ps *pairs) Push(x interface{}) {
	*ps = append(*ps, x.(pair))
}

func (ps *pairs) Pop() interface{} {
	x := (*ps)[ps.Len()-1]
	*ps = (*ps)[:ps.Len()-1]
	return x
}

func kSmallestPairs(nums1 []int, nums2 []int, k int) (res [][]int) {
	m, n := len(nums1), len(nums2)
	if m == 0 || n == 0 {
		return
	}
	idx := make([]int, m)
	if m*n < k {
		k = m * n
	}
	var ps pairs = make([]pair, 0)
	for i := 0; i < m; i++ {
		heap.Push(&ps, pair{nums1[i] + nums2[0], i})
	}
	for i := 0; i < k; i++ {
		x := heap.Pop(&ps).(pair)
		res = append(res, []int{nums1[x._2], nums2[idx[x._2]]})
		idx[x._2]++
		if idx[x._2] < n {
			heap.Push(&ps, pair{nums1[x._2] + nums2[idx[x._2]], x._2})
		}
	}
	return
}

type nx2Arr [][]int

func (arr nx2Arr) Len() int           { return len(arr) }
func (arr nx2Arr) Less(i, j int) bool { return arr[i][0]+arr[i][1] < arr[j][0]+arr[j][1] }
func (arr nx2Arr) Swap(i, j int)      { arr[i], arr[j] = arr[j], arr[i] }

func kSmallestPairsBS(nums1 []int, nums2 []int, k int) [][]int { // Brute search
	m, n := len(nums1), len(nums2)
	var sorted nx2Arr = make([][]int, m*n)
	for l, i := 0, 0; i < m; i++ {
		for j := 0; j < n; j++ {
			sorted[l] = []int{nums1[i], nums2[j]}
			l++
		}
	}
	sort.Sort(sorted)
	if m*n <= k {
		return sorted
	}
	return sorted[:k]
}