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:
-
Initialize the min-heap
q
. We start by forming pairs by taking each element fromnums1
(up to the firstk
elements to avoid unnecessary work) and pairing it with the first element innums2
. The reason for pairing with the first element fromnums2
is that we are trying to create pairs with the smallest possible sum to begin with. -
Heapify
q
to ensure the smallest sum pair is at the top of the heap. -
Prepare a list
ans
to store the resulting k smallest sum pairs. -
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 samenums1
element. If so, add this new pair to the heap so that it can be considered in the next iteration.
-
-
Return the
ans
list containing up tok
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.
-
Min-Heap Initialization and Heapify: The code snippet
q = [[u + nums2[0], i, 0] for i, u in enumerate(nums1[:k])]
initializes our min-heapq
. For each element innums1
(limited to the firstk
elements), we create a list with the sum (u + nums2[0]
), the index of thenums1
element (i
), and the index0
indicating we've used the first element ofnums2
. Theheapify(q)
call then converts the listq
into a heap data structure. -
Processing the Min-Heap: The
while
loop continues as long as there are elements in the heap (i.e., potential pairs to consider) andk
is greater than 0. In each iteration, we useheappop(q)
to remove and return the smallest element from the heap, which is the pair with the current smallest sum. -
Recording the Result: The current smallest sum pair's indices in
nums1
andnums2
are stored in the_
,i
,j
variables. We then append this pair[nums1[i], nums2[j]]
to our answer listans
. -
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 samenums1
element. Theif j + 1 < len(nums2)
check is used for this purpose, and if true, we form a new pair with the same element fromnums1
and the next element innums2
, then push it into the heap usingheappush(q, [nums1[i] + nums2[j + 1], i, j + 1])
to make sure it can be considered in subsequent iterations. -
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 EvaluatorExample 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:
-
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 firstk
pairs, andk
is3
, we consider only the first three elements fromnums1
. Therefore, our min-heapq
starts with the following elements (assuming we include all three elements ofnums1
andk
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. -
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. -
Recording the Result: We extract the indices
0
fromnums1
and0
fromnums2
and append the pair[1, 2]
to our result listans
. Now,ans = [[1, 2]]
. -
Generating New Pairs: Since
nums1[0]
was paired withnums2[0]
, we now pairnums1[0]
with the next element innums2
, which isnums2[1]
(value4
), and add this new pair to the heap. Now, the heapq
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)
. -
Finding the k Smallest Sums: We continue the process, now popping
(5, 0, 1)
from the heap and appending[1, 4]
toans
. Subsequently, we form a new pair withnums1[0]
andnums2[2]
, and the heapq
gets updated after heapification. We repeat this process untilk
becomes0
.
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:
-
Heap Initialization: The code creates a min-heap and initializes it with the sums of elements from
nums1
and the first element ofnums2
. Since the heap is initialized with at mostk
elements (and not more than the length ofnums1
), the complexity of this step isO(k)
since each heap insertion isO(log k)
and we do it at mostk
times. -
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 byk
and the length of the output), and bothheappop
andheappush
operations have a time complexity ofO(log k)
, the complexity for this part isO(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:
-
Heap Space: The heap size is at most
k
, as it stores pairs of indices and their corresponding sum, giving usO(k)
. -
Output List: The list
ans
to store the answer pairs. In the worst case, it will containk
pairs, leading toO(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.
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
https algomonster s3 us east 2 amazonaws com cover_photos heap svg Priority Queue and Heap What is the relationship between priority queue and heap Priority Queue is an Abstract Data Type and Heap is the concrete data structure we use to implement a priority queue Priority Queue A priority queue
LeetCode Patterns Your Personal Dijkstra's Algorithm to Landing Your Dream Job The goal of AlgoMonster is to help you get a job in the shortest amount of time possible in a data driven way We compiled datasets of tech interview problems and broke them down by patterns This way we
Recursion Recursion is one of the most important concepts in computer science Simply speaking recursion is the process of a function calling itself Using a real life analogy imagine a scenario where you invite your friends to lunch https algomonster s3 us east 2 amazonaws com recursion jpg You first
Want a Structured Path to Master System Design Too? Don’t Miss This!