Array
Heap (Priority Queue)
Description
You are given two integer arrays nums1
and nums2
sorted in non-decreasing 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 (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]
Constraints:
1 <= nums1.length, nums2.length <= 105
-109 <= nums1[i], nums2[i] <= 109
nums1
and nums2
both are sorted in non-decreasing order .
1 <= k <= 104
k <= nums1.length * nums2.length
Solutions
Solution 1
Python3 Java C++ Go
1
2
3
4
5
6
7
8
9
10
11
12
13
14 class Solution :
def kSmallestPairs (
self , nums1 : List [ int ], nums2 : List [ int ], k : int
) -> List [ List [ int ]]:
q = [[ u + nums2 [ 0 ], i , 0 ] for i , u in enumerate ( nums1 [: k ])]
heapify ( q )
ans = []
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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18 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 ;
}
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24 class Solution {
public :
vector < vector < int >> kSmallestPairs ( vector < int >& nums1 , vector < int >& nums2 , int k ) {
auto cmp = [ & nums1 , & nums2 ]( const pair < int , int >& a , const pair < int , int >& b ) {
return nums1 [ a . first ] + nums2 [ a . second ] > nums1 [ b . first ] + nums2 [ b . second ];
};
int m = nums1 . size ();
int n = nums2 . size ();
vector < vector < int >> ans ;
priority_queue < pair < int , int > , vector < pair < int , int >> , decltype ( cmp ) > pq ( cmp );
for ( int i = 0 ; i < min ( k , m ); i ++ )
pq . emplace ( i , 0 );
while ( k -- && ! pq . empty ()) {
auto [ x , y ] = pq . top ();
pq . pop ();
ans . emplace_back ( initializer_list < int > { nums1 [ x ], nums2 [ y ]});
if ( y + 1 < n )
pq . emplace ( x , y + 1 );
}
return ans ;
}
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31 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 any ) { h . data = append ( h . data , v .( pair )) }
func ( h * hp ) Pop () any { a := h . data ; v := a [ len ( a ) - 1 ]; h . data = a [: len ( a ) - 1 ]; return v }
GitHub