373. Find K Pairs with Smallest Sums

Problem

You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k.

Define a pair (u,v) which consists of one element from the first array and one element from the second array.

Find the k pairs (u1,v1),(u2,v2) ...(uk,vk) with the smallest sums.

Example 1:

Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]] 
Explanation: The first 3 pairs are returned from the sequence: 
             [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

Example 2:

Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [1,1],[1,1]
Explanation: The first 2 pairs are returned from the sequence: 
             [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

Example 3:

Input: nums1 = [1,2], nums2 = [3], k = 3
Output: [1,3],[2,3]
Explanation: All possible pairs are returned from the sequence: [1,3],[2,3]

Problem Summary

Given two integer arrays nums1 and nums2 sorted in ascending order, and an integer k.

Define a pair (u,v), where the first element comes from nums1 and the second element comes from nums2.

Find the k pairs of numbers (u1,v1), (u2,v2) ... (uk,vk) with the smallest sums.

Solution Ideas

• Given 2 arrays and a number k, the task is to find k pairs of values whose two numbers have the smallest sums.
• At first glance, this problem seems solvable with binary search. The two arrays have m * n possible pairs. Then find the minimum sum and the maximum sum, and perform binary search within this range. For each mid, count how many pairs have sums smaller than mid. If the count is less than k, continue binary searching in the right interval; if the count is greater than k, continue binary searching in the left interval. So far, this idea seems feasible. However, the pairs found during each search are unordered. This can lead to incorrect final results. For example, when mid = 10, there are 22 sums less than 10, while k = 25. This means mid is too small, so increase mid. When mid = 11, there are 30 sums less than 11, while k = 25. At this point, we should take the first 25 from these 30 sums. But when we traverse the pairs, the sums are not sorted from smallest to largest. So we still need to additionally sort these 30 candidate values. This increases the time complexity again.
• We can first solve it with a brute-force approach. Traverse all sums, sort them, and take the first k. This brute-force method can be accepted.
• The optimal solution for this problem should be a priority queue. Maintain a min-heap. Put the sums of the pairs into this min-heap, and continuously pop the k smallest values into an array, which is the answer.
• This series of problems about finding the K-th smallest element in a sorted matrix includes: Problem 373, Problem 378, Problem 668, Problem 719, and Problem 786.

Code

package leetcode

import (
	"container/heap"
	"sort"
)

// Solution 1: Priority queue
func kSmallestPairs(nums1 []int, nums2 []int, k int) [][]int {
	result, h := [][]int{}, &minHeap{}
	if len(nums1) == 0 || len(nums2) == 0 || k == 0 {
		return result
	}
	if len(nums1)*len(nums2) < k {
		k = len(nums1) * len(nums2)
	}
	heap.Init(h)
	for _, num := range nums1 {
		heap.Push(h, []int{num, nums2[0], 0})
	}
	for len(result) < k {
		min := heap.Pop(h).([]int)
		result = append(result, min[:2])
		if min[2] < len(nums2)-1 {
			heap.Push(h, []int{min[0], nums2[min[2]+1], min[2] + 1})
		}
	}
	return result
}

type minHeap [][]int

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

func (h *minHeap) Push(x interface{}) {
	*h = append(*h, x.([]int))
}

func (h *minHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// Solution 2: Brute force
func kSmallestPairs1(nums1 []int, nums2 []int, k int) [][]int {
	size1, size2, res := len(nums1), len(nums2), [][]int{}
	if size1 == 0 || size2 == 0 || k < 0 {
		return nil
	}
	for i := 0; i < size1; i++ {
		for j := 0; j < size2; j++ {
			res = append(res, []int{nums1[i], nums2[j]})
		}
	}
	sort.Slice(res, func(i, j int) bool {
		return res[i][0]+res[i][1] < res[j][0]+res[j][1]
	})
	if len(res) >= k {
		return res[:k]
	}
	return res
}