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] }