373. Find K Pairs with Smallest Sums


Problem Description

In this problem, we are working with two key concepts: pairs and sum minimization. We have two arrays, nums1 and nums2, both sorted in non-decreasing order. Our task is to find k pairs, each composed of one element from each array. The objective is to identify the pairs with the smallest possible sums. Specifically, we want to find the k pairs whose sums are the smallest among all possible combinations from nums1 and nums2.

If we were to list all possible pairs and their sums, the list would be quite extensive. However, since the arrays are already sorted, we can infer that the smallest sums will involve the smallest elements from both arrays.

Intuition

The solution employs a min-heap, a binary tree where the parent node is always less than or equal to its children. This is an ideal data structure for efficiently finding the smallest elements.

The algorithm initializes by creating a min-heap and pre-filling it with pairs combining each of the first k elements of nums1 with the first element of nums2.

Here's the step-by-step process to understand the solution:

  1. Initialize the min-heap q. We start by forming pairs by taking each element from nums1 (up to the first k elements to avoid unnecessary work) and pairing it with the first element in nums2. The reason for pairing with the first element from nums2 is that we are trying to create pairs with the smallest possible sum to begin with.

  2. Heapify q to ensure the smallest sum pair is at the top of the heap.

  3. Prepare a list ans to store the resulting k smallest sum pairs.

  4. Use a loop to process the min-heap. While the min-heap is not empty and k is greater than 0, we do the following:

    • Pop the smallest sum pair from the heap. This pair is guaranteed to be one of the smallest sum pairs because of the min-heap's properties.

    • Append this pair to the ans list.

    • Decrease k by 1, as we have successfully found one of the k smallest sum pairs.

    • Check if there's a next element in nums2 that can be paired with the same nums1 element. If so, add this new pair to the heap so that it can be considered in the next iteration.

  5. Return the ans list containing up to k smallest sum pairs.

By utilizing a min-heap, we ensure that at each step, we pop the minimum sum pair without having to check all combinations. Moreover, since nums1 and nums2 are sorted, we can confidently move to the next element in nums2 for potential pairing without worrying about missing smaller sums.

Learn more about Heap (Priority Queue) patterns.

Solution Approach

The solution we're discussing uses a min-heap to efficiently find the k pairs with the smallest sums. Let's walk through the implementation step by step, taking advantage of algorithms, data structures, and patterns.

  1. Min-Heap Initialization and Heapify: The code snippet q = [[u + nums2[0], i, 0] for i, u in enumerate(nums1[:k])] initializes our min-heap q. For each element in nums1 (limited to the first k elements), we create a list with the sum (u + nums2[0]), the index of the nums1 element (i), and the index 0 indicating we've used the first element of nums2. The heapify(q) call then converts the list q into a heap data structure.

  2. Processing the Min-Heap: The while loop continues as long as there are elements in the heap (i.e., potential pairs to consider) and k is greater than 0. In each iteration, we use heappop(q) to remove and return the smallest element from the heap, which is the pair with the current smallest sum.

  3. Recording the Result: The current smallest sum pair's indices in nums1 and nums2 are stored in the _, i, j variables. We then append this pair [nums1[i], nums2[j]] to our answer list ans.

  4. Generating New Pairs: After finding a pair with a small sum, we check if there is a next element in nums2 to pair with the same nums1 element. The if j + 1 < len(nums2) check is used for this purpose, and if true, we form a new pair with the same element from nums1 and the next element in nums2, then push it into the heap using heappush(q, [nums1[i] + nums2[j + 1], i, j + 1]) to make sure it can be considered in subsequent iterations.

  5. Finding the k Smallest Sums: The heap keeps track of the current smallest pairs, and the loop ensures that each time we pop from the heap and push a new pair into it, we're considering the next smallest possible pair. This continues until we've found k pairs or there are no more new pairs to consider.

The key algorithms and data structures used here are the heap data structure (a min-heap in particular) for keeping track of the smallest pairs and a list for our result.

The pattern is utilizing a greedy-like approach: in each step, we're always taking the smallest available option (the pair with the smallest sum) and then exploring slightly larger options by considering the next possible element from nums2.

This approach results in a time-efficient solution because we avoid generating all possible pairs, which would have resulted in a much higher time complexity.

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's walk through a small example using the provided solution approach to understand how it works.

Consider the following input:

  • nums1 = [1, 7, 11]
  • nums2 = [2, 4, 6]
  • k = 3

Following the steps of the solution approach:

  1. Min-Heap Initialization and Heapify: We initialize the min-heap with elements paired as follows: (nums1[i] + nums2[0], i, 0). Since we're only interested in the first k pairs, and k is 3, we consider only the first three elements from nums1. Therefore, our min-heap q starts with the following elements (assuming we include all three elements of nums1 and k limits other factors):

    q = [
      (1 + 2, 0, 0),   # sum = 3, index in nums1 = 0, index in nums2 = 0
      (7 + 2, 1, 0),   # sum = 9, index in nums1 = 1, index in nums2 = 0
      (11 + 2, 2, 0),  # sum = 13, index in nums1 = 2, index in nums2 = 0
    ]

    After heapification, q ensures that the pair with the smallest sum is at the top of the heap.

  2. Processing the Min-Heap: We start processing the heap by popping elements while k > 0. The smallest sum pair is at the top, so we pop (3, 0, 0) first.

  3. Recording the Result: We extract the indices 0 from nums1 and 0 from nums2 and append the pair [1, 2] to our result list ans. Now, ans = [[1, 2]].

  4. Generating New Pairs: Since nums1[0] was paired with nums2[0], we now pair nums1[0] with the next element in nums2, which is nums2[1] (value 4), and add this new pair to the heap. Now, the heap q looks like:

    q = [
      (1 + 4, 0, 1),   # sum = 5, new pair
      (7 + 2, 1, 0),
      (11 + 2, 2, 0)
    ]

    Heapification ensures the smallest pair stays at the top, which is now (5, 0, 1).

  5. Finding the k Smallest Sums: We continue the process, now popping (5, 0, 1) from the heap and appending [1, 4] to ans. Subsequently, we form a new pair with nums1[0] and nums2[2], and the heap q gets updated after heapification. We repeat this process until k becomes 0.

After these steps, our result list ans is [[1, 2], [1, 4], [1, 6]], which are the k smallest sum pairs from nums1 and nums2.

By following this strategy, we avoid generating all possible pairs and efficiently find the k smallest sums using a min-heap to guide our pair selections.

Solution Implementation

1from typing import List
2from heapq import heapify, heappop, heappush
3
4class Solution:
5    def k_smallest_pairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
6        # Create a priority queue to hold the sum of pairs along with the indices in nums1 and nums2
7        # Only consider the first k numbers in nums1 for initial pairing, as we are looking for the k smallest pairs
8        priority_queue = [[u + nums2[0], i, 0] for i, u in enumerate(nums1[:k])]
9        heapify(priority_queue)  # Convert list into a heap
10
11        result = []  # Initialize a list to hold the result pairs
12
13        # Iterate until the priority queue is empty or we have found k pairs
14        while priority_queue and k > 0:
15            # Pop the smallest sum pair from the heap
16            sum_pair, index1, index2 = heappop(priority_queue)
17          
18            # Append the corresponding values from nums1 and nums2 to the result
19            result.append([nums1[index1], nums2[index2]])
20          
21            k -= 1  # Decrement k as we have found one of the k smallest pairs
22          
23            # If there are more elements in nums2 to pair with the current nums1 element, push the next pair onto the heap
24            if index2 + 1 < len(nums2):
25                heappush(priority_queue, [nums1[index1] + nums2[index2 + 1], index1, index2 + 1])
26
27        return result  # Return the list of k smallest pairs
28
1class Solution {
2    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
3        // Create a priority queue to hold the arrays with a comparator to prioritize by the sum of pairs.
4        PriorityQueue<int[]> queue = new PriorityQueue<>((a, b) -> a[0] - b[0]);
5
6        // Initialize the priority queue with the first k pairs from nums1 and the first element from nums2.
7        for (int i = 0; i < Math.min(nums1.length, k); ++i) {
8            queue.offer(new int[] {nums1[i] + nums2[0], i, 0});
9        }
10      
11        // Prepare a list to store the k smallest pairs.
12        List<List<Integer>> result = new ArrayList<>();
13      
14        // Keep polling from the priority queue to find the next smallest pair
15        while (!queue.isEmpty() && k > 0) {
16            // Poll the smallest sum pair from the priority queue.
17            int[] currentPair = queue.poll();
18          
19            // Add the new pair [nums1[index1], nums2[index2]] to the result list.
20            result.add(Arrays.asList(nums1[currentPair[1]], nums2[currentPair[2]]));
21          
22            // Decrease the remaining pairs count.
23            --k;
24          
25            // If there's a next element in nums2, offer the next pair from nums1 and nums2 into the priority queue.
26            if (currentPair[2] + 1 < nums2.length) {
27                queue.offer(new int[] {nums1[currentPair[1]] + nums2[currentPair[2] + 1], currentPair[1], currentPair[2] + 1});
28            }
29        }
30      
31        // Return the list of k smallest pairs.
32        return result;
33    }
34}
35
1#include <vector>
2#include <queue>
3using namespace std;
4
5class Solution {
6public:
7    // Function to find k smallest pairs with minimal sum
8    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
9        // Lambda function to compare pairs based on the sum of elements they point to in nums1 and nums2
10        auto comparePairs = [&nums1, &nums2](const pair<int, int>& a, const pair<int, int>& b) {
11            return nums1[a.first] + nums2[a.second] > nums1[b.first] + nums2[b.second];
12        };
13
14        // Get the sizes of the input arrays
15        int size1 = nums1.size();
16        int size2 = nums2.size();
17
18        // Variable to store the final pairs
19        vector<vector<int>> result;
20      
21        // Priority queue to keep pairs in ascending order based on their sum
22        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(comparePairs)> minHeap(comparePairs);
23      
24        // Initialize the priority queue with the first pair from each element in nums1
25        for (int i = 0; i < min(k, size1); ++i) {
26            minHeap.emplace(i, 0);
27        }
28
29        // Extract the k smallest pairs
30        while (k-- > 0 && !minHeap.empty()) {
31            auto [index1, index2] = minHeap.top();
32            minHeap.pop();
33
34            // Add the current smallest pair to the result
35            result.push_back({nums1[index1], nums2[index2]});
36          
37            // If there's a next element in nums2, add the new pair to the priority queue
38            if (index2 + 1 < size2) {
39                minHeap.emplace(index1, index2 + 1);
40            }
41        }
42
43        // Return all k smallest pairs found
44        return result;
45    }
46};
47
1// Importing required modules
2import { PriorityQueue } from 'typescript-collections'; // Placeholder import for PQ functionality
3
4// Defining the type for a pair of indices
5type IndexPair = [number, number];
6
7// Function to compare pairs based on the sum of elements they point to in nums1 and nums2
8function comparePairs(nums1: number[], nums2: number[], a: IndexPair, b: IndexPair): boolean {
9  return nums1[a[0]] + nums2[a[1]] > nums1[b[0]] + nums2[b[1]];
10}
11
12// Function to find k smallest pairs with minimal sum
13function kSmallestPairs(nums1: number[], nums2: number[], k: number): number[][] {
14  // Get the sizes of the input arrays
15  const size1 = nums1.length;
16  const size2 = nums2.length;
17
18  // Variable to store the final pairs
19  const result: number[][] = [];
20
21  // Priority queue to keep pairs in ascending order based on their sum
22  const minHeap = new PriorityQueue<IndexPair>((a, b) => comparePairs(nums1, nums2, a, b));
23
24  // Initialize the priority queue with the first pair from each element in nums1
25  for (let i = 0; i < Math.min(k, size1); i++) {
26    minHeap.enqueue([i, 0]);
27  }
28
29  // Extract the k smallest pairs
30  while (k > 0 && !minHeap.isEmpty()) {
31    const [index1, index2] = minHeap.dequeue()!;
32
33    // Add the current smallest pair to the result
34    result.push([nums1[index1], nums2[index2]]);
35  
36    // If there's a next element in nums2, add the new pair to the priority queue
37    if (index2 + 1 < size2) {
38      minHeap.enqueue([index1, index2 + 1]);
39    }
40
41    k--;
42  }
43
44  // Return all k smallest pairs found
45  return result;
46}
47

Time and Space Complexity

The given Python code implements a heap to find the k smallest pairs of sums from two integer arrays nums1 and nums2. Here we will discuss its time complexity and space complexity.

Time Complexity

The time complexity of the algorithm is as follows:

  1. Heap Initialization: The code creates a min-heap and initializes it with the sums of elements from nums1 and the first element of nums2. Since the heap is initialized with at most k elements (and not more than the length of nums1), the complexity of this step is O(k) since each heap insertion is O(log k) and we do it at most k times.

  2. Heap Operations: Then, in each iteration, the algorithm pops an element from the heap and potentially pushes a new element into the heap. Since we perform k iterations (bounded by k and the length of the output), and both heappop and heappush operations have a time complexity of O(log k), the complexity for this part is O(k log k).

Overall, the time complexity is O(k) + O(k log k), which simplifies to O(k log k) because as k grows, the k log k term dominates.

Space Complexity

The space complexity of the algorithm is as follows:

  1. Heap Space: The heap size is at most k, as it stores pairs of indices and their corresponding sum, giving us O(k).

  2. Output List: The list ans to store the answer pairs. In the worst case, it will contain k pairs, leading to O(k) space complexity.

The overall space complexity combines the heap space and the output list, but since both are O(k), the total space complexity remains O(k).

Learn more about how to find time and space complexity quickly using problem constraints.


Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:
Question 1 out of 10

What's the output of running the following function using input [30, 20, 10, 100, 33, 12]?

1def fun(arr: List[int]) -> List[int]:
2    import heapq
3    heapq.heapify(arr)
4    res = []
5    for i in range(3):
6        res.append(heapq.heappop(arr))
7    return res
8
1public static int[] fun(int[] arr) {
2    int[] res = new int[3];
3    PriorityQueue<Integer> heap = new PriorityQueue<>();
4    for (int i = 0; i < arr.length; i++) {
5        heap.add(arr[i]);
6    }
7    for (int i = 0; i < 3; i++) {
8        res[i] = heap.poll();
9    }
10    return res;
11}
12
1class HeapItem {
2    constructor(item, priority = item) {
3        this.item = item;
4        this.priority = priority;
5    }
6}
7
8class MinHeap {
9    constructor() {
10        this.heap = [];
11    }
12
13    push(node) {
14        // insert the new node at the end of the heap array
15        this.heap.push(node);
16        // find the correct position for the new node
17        this.bubble_up();
18    }
19
20    bubble_up() {
21        let index = this.heap.length - 1;
22
23        while (index > 0) {
24            const element = this.heap[index];
25            const parentIndex = Math.floor((index - 1) / 2);
26            const parent = this.heap[parentIndex];
27
28            if (parent.priority <= element.priority) break;
29            // if the parent is bigger than the child then swap the parent and child
30            this.heap[index] = parent;
31            this.heap[parentIndex] = element;
32            index = parentIndex;
33        }
34    }
35
36    pop() {
37        const min = this.heap[0];
38        this.heap[0] = this.heap[this.size() - 1];
39        this.heap.pop();
40        this.bubble_down();
41        return min;
42    }
43
44    bubble_down() {
45        let index = 0;
46        let min = index;
47        const n = this.heap.length;
48
49        while (index < n) {
50            const left = 2 * index + 1;
51            const right = left + 1;
52
53            if (left < n && this.heap[left].priority < this.heap[min].priority) {
54                min = left;
55            }
56            if (right < n && this.heap[right].priority < this.heap[min].priority) {
57                min = right;
58            }
59            if (min === index) break;
60            [this.heap[min], this.heap[index]] = [this.heap[index], this.heap[min]];
61            index = min;
62        }
63    }
64
65    peek() {
66        return this.heap[0];
67    }
68
69    size() {
70        return this.heap.length;
71    }
72}
73
74function fun(arr) {
75    const heap = new MinHeap();
76    for (const x of arr) {
77        heap.push(new HeapItem(x));
78    }
79    const res = [];
80    for (let i = 0; i < 3; i++) {
81        res.push(heap.pop().item);
82    }
83    return res;
84}
85

Recommended Readings

Want a Structured Path to Master System Design Too? Don’t Miss This!