[LeetCode] Challenge log 373

Posted on | Edited on |
373. Find K Pairs with Smallest Sums

Tried many methods. Take it as a review to heap and piority queue…

  1. qsort —> time limit exceed
  2. build min heap, pop the first k pairs —> ac, beat 9%
  3. use max heap to maintain the k min pairs —> ac, beat 14%
  4. use the fact that both arrays are presorted, and piority queue to maintain the candidates —> ac, beat 52%
  5. use built-in library “heapq” —> ac, beat 74%
  6. use set, instead of array to mark visted slots —> ac, beat 95%

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:

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

Return: [1,2],[1,4],[1,6]

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:

1
2
3
4
5
6
Given nums1 = [1,1,2], nums2 = [1,2,3],  k = 2

Return: [1,1],[1,1]

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:

1
2
3
4
5
6
Given nums1 = [1,2], nums2 = [3],  k = 3 

Return: [1,3],[2,3]

All possible pairs are returned from the sequence:
[1,3],[2,3]

Credits:
Special thanks to @elmirap and @StefanPochmann for adding this problem and creating all test cases.

Soulution:

First we can visualize the all pairs as below:

0 1 2 3
0 added added added candidate
1 added candidate
2 candidate

We observes that, for (i, j), only when (i-1, j) and (i, j-1) are added then it can be a candidate. With this lead, my solution keeps a piority queue, from which the min pair is popped out and then added to the ouput each time. Meanwhile we check if (i, j+1) or (i+1, j) is a valid candidate. If so, add it to the queue. Repeat until ouput has k pairs.

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
from heapq import *

class Solution:
    def kSmallestPairs(self, nums1, nums2, k):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :type k: int
        :rtype: List[List[int]]
        """
 
        ncol = min(k, len(nums2))
        nrow = min(k, len(nums1))
        if ncol == 0 or nrow == 0:return []
        
        mark = set()
        cand = [(nums1[0] + nums2[0] ,0,0)]
        res = []
        while len(res) < k and len(cand) > 0:
            s, i, j = heappop(cand)
            res.append([nums1[i], nums2[j]])
            mark.add((i, j))
            if i+1 < nrow and ((j-1>=0 and (i+1, j-1) not in mark) or (j-1 < 0)):
                heappush(cand, (nums1[i+1] + nums2[j], i+1, j))
            if j+1 < ncol and ((i-1>=0 and (i-1, j+1) not in mark) or (i-1 < 0)):
                heappush(cand, (nums1[i] + nums2[j+1], i, j+1))
        return res
0%