Question

Formatted question description: https://leetcode.ca/all/373.html

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.

Return the k pairs (u₁, v₁), (u₂, v₂), ..., (u_k, v_k) 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]

Constraints:

1 <= nums1.length, nums2.length <= 10⁵
-10⁹ <= nums1[i], nums2[i] <= 10⁹
nums1 and nums2 both are sorted in ascending order.
1 <= k <= 10⁴

Algorithm

The value of num1[0]+num2[0] must be the smallest, so we first throw this value into the heap.

Then start traversing the next element.

The next value must be the smallest one from num1[1]+num2[0] and num1[0]+num2[1].

Loop all the combinations in turn until the extracted element is K.

Code

import java.util.*;

public class Find_K_Pairs_with_Smallest_Sums {

    public static void main(String[] args) {
        Find_K_Pairs_with_Smallest_Sums out = new Find_K_Pairs_with_Smallest_Sums();
        Solution s = out.new Solution();

        // output: [[1,2],[1,4],[1,6]]
//        s.kSmallestPairs(new int[]{1,7,11}, new int[]{2,4,6}, 3).forEach(e -> System.out.println(e.toString()));

        s.kSmallestPairs(new int[]{1,1,2}, new int[]{1,2,3}, 2).forEach(e -> System.out.println(e.toString()));
    }

    public class Solution_NlogN {
        public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {

            List<List<Integer>> result = new ArrayList<>();

            if(nums1.length == 0 || nums2.length == 0 || k == 0) {
                return result;
            }

            PriorityQueue<int[]> q = new PriorityQueue<>(
                (a, b) -> a[0]+a[1]-b[0]-b[1]
            ); // @note: min at top of heap

            // all nums-1 to in heap
            for(int i = 0; i < nums1.length && i < k; i++) {

                // last 0 meaning index of 'nums2', to check if this index is hitting nums end
                int[] oneVector = new int[]{nums1[i], nums2[0], 0};
                q.offer(oneVector);
            }

            while(k > 0 && !q.isEmpty()){

                k--; // not end of loop block, so that below continue can safely run to exit while

                int[] pair = q.poll();
                result.add(Arrays.asList(pair[0], pair[1]));

                if(pair[2] == nums2.length-1) {
                    continue;
                }

                int nextIndexInNums1 = pair[0];
                int nextIndexInNums2 = pair[2] + 1;
                q.offer(new int[]{nextIndexInNums1, nums2[nextIndexInNums2], nextIndexInNums2});
            }

            return result;
        }
    }

    class Solution {
        public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {

            List<List<Integer>> result = new ArrayList<>();

            if(nums1.length == 0 || nums2.length == 0){
                return result;
            }

            PriorityQueue<List<Integer>> queue = new PriorityQueue<>(
                k, // initial size will not block auto-grow
                (a,b) -> b.get(0) + b.get(1) - a.get(0) - a.get(1)
            );

            for (int i = 0; i < Math.min(k, nums1.length); i++) {
                for (int j = 0; j < Math.min(k, nums2.length); j++) {
                    if (queue.size() < k) {
                        queue.offer(Arrays.asList(nums1[i], nums2[j]));
                    } else if (nums1[i] + nums2[j] < queue.peek().get(0) + queue.peek().get(1)) {
                        queue.offer(Arrays.asList(nums1[i], nums2[j]));
                        queue.poll();
                    }
                }
            }

            return new ArrayList<>(queue);
        }
    }
}

############

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<int[]> q = new PriorityQueue<>(Comparator.comparingInt(a -> a[0]));
        for (int i = 0; i < Math.min(nums1.length, k); ++i) {
            q.offer(new int[] {nums1[i] + nums2[0], i, 0});
        }
        List<List<Integer>> ans = new ArrayList<>();
        while (!q.isEmpty() && k > 0) {
            int[] e = q.poll();
            ans.add(Arrays.asList(nums1[e[1]], nums2[e[2]]));
            --k;
            if (e[2] + 1 < nums2.length) {
                q.offer(new int[] {nums1[e[1]] + nums2[e[2] + 1], e[1], e[2] + 1});
            }
        }
        return ans;
    }
}

class Solution {
public:
    vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<pair<int, int>> res;
        priority_queue<pair<int, int>, vector<pair<int, int>>, cmp> q;
        for (int i = 0; i < min((int)nums1.size(), k); ++i) {
            for (int j = 0; j < min((int)nums2.size(), k); ++j) {
                if (q.size() < k) {
                    q.push({nums1[i], nums2[j]});
                } else if (nums1[i] + nums2[j] < q.top().first + q.top().second) {
                    q.push({nums1[i], nums2[j]}); q.pop();
                }
            }
        }
        while (!q.empty()) {
            res.push_back(q.top()); q.pop();
        }
        return res;
    }
    struct cmp {
        bool operator() (pair<int, int> &a, pair<int, int> &b) {
            return a.first + a.second < b.first + b.second;
        }
    };
};

class Solution:
    def kSmallestPairs(
        self, nums1: List[int], nums2: List[int], k: int
    ) -> List[List[int]]:
        '''
            >>> [1,2,3][:99]
            [1, 2, 3]
        '''
        q = [[u + nums2[0], i, 0] for i, u in enumerate(nums1[:k])] # still need '[u + nums2[0]', for q ordering
        heapify(q)
        ans = []
        # both q and k should be checked
        # because k can be super larger than nums1+nums2
        while q and k > 0:
            _, i, j = heappop(q)
            ans.append([nums1[i], nums2[j]])
            k -= 1
            if j + 1 < len(nums2):
                heappush(q, [nums1[i] + nums2[j + 1], i, j + 1])
        return ans

############

import heapq


class Solution(object):
  def kSmallestPairs(self, nums1, nums2, k):
    """
    :type nums1: List[int]
    :type nums2: List[int]
    :type k: int
    :rtype: List[List[int]]
    """
    if not nums1 or not nums2:
      return []
    heap = [(nums1[0] + nums2[0], 0, 0)]
    ans = []
    visited = {(0, 0)}

    while heap:
      val, i, j = heapq.heappop(heap)
      ans.append((nums1[i], nums2[j]))
      k -= 1
      if k == 0:
        return ans
      if i + 1 < len(nums1) and (i + 1, j) not in visited:
        heapq.heappush(heap, (nums1[i + 1] + nums2[j], i + 1, j))
        visited |= {(i + 1, j)}
      if j + 1 < len(nums2) and (i, j + 1) not in visited:
        heapq.heappush(heap, (nums1[i] + nums2[j + 1], i, j + 1))
        visited |= {(i, j + 1)}
    return ans

func kSmallestPairs(nums1, nums2 []int, k int) (ans [][]int) {
	m, n := len(nums1), len(nums2)
	h := hp{nil, nums1, nums2}
	for i := 0; i < k && i < m; i++ {
		h.data = append(h.data, pair{i, 0})
	}
	for h.Len() > 0 && len(ans) < k {
		p := heap.Pop(&h).(pair)
		i, j := p.i, p.j
		ans = append(ans, []int{nums1[i], nums2[j]})
		if j+1 < n {
			heap.Push(&h, pair{i, j + 1})
		}
	}
	return
}

type pair struct{ i, j int }
type hp struct {
	data         []pair
	nums1, nums2 []int
}

func (h hp) Len() int { return len(h.data) }
func (h hp) Less(i, j int) bool {
	a, b := h.data[i], h.data[j]
	return h.nums1[a.i]+h.nums2[a.j] < h.nums1[b.i]+h.nums2[b.j]
}
func (h hp) Swap(i, j int)       { h.data[i], h.data[j] = h.data[j], h.data[i] }
func (h *hp) Push(v interface{}) { h.data = append(h.data, v.(pair)) }
func (h *hp) Pop() interface{}   { a := h.data; v := a[len(a)-1]; h.data = a[:len(a)-1]; return v }

373 - Find K Pairs with Smallest Sums

Question

Algorithm

Code

All Problems

All Solutions

Welcome to Subscribe On Youtube

Question

Algorithm

Code

All Problems

All Solutions