Skip to content

373. Find K Pairs with Smallest Sums

Problem

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 (u1, v1), (u2, v2), ..., (uk, vk) with the smallest sums.

Constraints:

  • 1 <= nums1.length, nums2.length <= 10**5
  • -10**9 <= nums1[i], nums2[i] <= 10**9
  • nums1 and nums2 both are sorted in ascending order.
  • 1 <= k <= 10**4

First look

  1. You want to get minimal value: Which could be heap, sorting, min binary tree,
  2. We don’t care about the duplicate or position so we could go for: heap, sorting
  3. Sorting give overall better performance, as the nums_1 and nums_2 is already sorted in ascending order. By define, sum of a pair number (x1,y1) will be less than (x2,y2) if x1.index <= x2.index and y1.index <= y2.index (x in nums_1, y in nums_2)

Solve

1. Quick first thought

  1. Create a pointer p for each num_arr , index from 0, we calling their created pair(p1, p2) which start at (0,0)
  2. Loop k times, push pair (p1, p2) into result arr and increase either p1 or p2 (base on value of nums_1[p1] and nums_2[p2])
  3. Return created arr

2. Wow, that wrong

  1. The first one don’t reuse index after p1 and p2 already passing it, so we can’t detect answer that use the index back to back
  2. Quick fix is running a max cap of p1 and p2, if we reach and break the max cap, we can reuse all element that in the range of either [0..p1] or [0..p2]

3. Heap attempt, update the pointer case

Seem like heap still coming back, as rerunning and pointer can’t handle all case where we actually need to consider the sum of pair (p1,p2)

  1. Quickly add all available pair create by nums1 and nums2 into each correspond separated heap, every heap element is a pair (p1, p2) with priority value of sum = nums1[p1] + nums2[p2]
  2. Update: Every element p1 of nums1 will be pair with the lowest available p2 of nums2 ; After append to the result; we will increase p2 to make the new pair from p1 and push it back to heap
  3. Run in a loop until we get k element in the result arr; We use a set() hash map to handle with the duplication

Implement

1. First implement

Quick and easy, follow first thought process

class Solution:
      def kSmallestPairs(self,
                         nums1: List[int],
                         nums2: List[int],
                         k: int) -> List[List[int]]:
          p1, p2 = 0, 0
          arr = []
          for _ in range(k):
              arr.append([nums1[p1], nums2[p2]])
              if not p1 + 1 < len(nums1) and not p2 + 1 < len(nums2):
                  break
              if not p1 + 1 < len(nums1):
                  p2 += 1
                  continue
              if not p2+1 < len(nums2):
                  p1 += 1
                  continue

              if nums1[p1+1] > nums2[p2+1]:
                  p2 += 1
              else:
                  p1 += 1

          return arr

2. Second attempt

After seeing the problem that come with reusing number

class Solution:
      def kSmallestPairs(self,
                         nums1: List[int],
                         nums2: List[int],
                         k: int) -> List[List[int]]:
          p1, p2 = 0, 0
          max_p1, max_p2 = p1, p2
          arr = []
          for _ in range(k):
              arr.append([nums1[p1], nums2[p2]])
              if not p1 + 1 < len(nums1) and not p2 + 1 < len(nums2):
                  break
              if not p1 + 1 < len(nums1):
                  p2 += 1
                  if p2 > max_p2:
                      max_p2 = p2
                      p1 = 0
                  continue
              if not p2+1 < len(nums2):
                  p1 += 1
                  if p1 > max_p1:
                      max_p1 = p1
                      p2 = 0
                  continue

              if nums1[p1+1] > nums2[p2+1]:
                  p2 += 1
                  if p2 > max_p2:
                      max_p2 = p2
                      p1 = 0
              else:
                  p1 += 1
                  if p1 > max_p1:
                      max_p1 = p1
                      p2 = 0

          return arr

3. Heap attempt

  • Running pointer for every element of both nums1 and nums2 instead
  • If we add a (pos1, pos2) pair into result arr; Try increase the correspond pointer, and push newly created pair either (pos1 +1, pos2) or (pos1 , pos2+1) to possible answer.
  • Handle the find minimal answer by using heap
  • Handle possible duplication by using a set() hash map
class Solution:
    def kSmallestPairs(self,
                       nums1: List[int],
                       nums2: List[int],
                       k: int) -> List[List[int]]:
        trace = set()
        p1 = [(value + nums2[0], (index, 0))
              for index, value in enumerate(nums1)]
        p2 = [(value + nums1[0], (0, index))
              for index, value in enumerate(nums2)]
        heapq.heapify(p1)
        heapq.heapify(p2)
        arr = []
        for _ in range(k*len(nums1)*len(nums2)):
            if len(trace) == k:
                break
            if len(p1) == len(p2) == 0:
                break

            needPopP1 = False
            if len(p1) == 0:
                needPopP1 = False
            elif len(p2) == 0:
                needPopP1 = True
            else:
                v1, (_, _) = p1[0]
                v2, (_, _) = p2[0]
                needPopP1 = v1 <= v2
            if needPopP1:
                _, (pos1, pos2) = heapq.heappop(p1)
                if (pos1, pos2) not in trace:
                    arr.append([nums1[pos1], nums2[pos2]])
                    trace.add((pos1, pos2))
                if pos2+1 < len(nums2):
                    heapq.heappush(
                        p1, (nums1[pos1] + nums2[pos2+1], (pos1, pos2+1)))
            else:
                _, (pos1, pos2) = heapq.heappop(p2)
                if (pos1, pos2) not in trace:
                    arr.append([nums1[pos1], nums2[pos2]])
                    trace.add((pos1, pos2))
                if pos1+1 < len(nums1):
                    heapq.heappush(
                        p2, (nums1[pos1+1] + nums2[pos2], (pos1+1, pos2)))

        return arr

4. Final refactor

  • The heap for nums2 isn’t necessary, we only need allocate pointer for every element of either nums1 or nums2 (This implement using nums1)
  • Also, this implementation won’t have duplicate value, so set() hash map isn’t necessary
class Solution:
    def kSmallestPairs(self,
                       nums1: List[int],
                       nums2: List[int],
                       k: int) -> List[List[int]]:
        p1 = [(value + nums2[0], (index, 0))
              for index, value in enumerate(nums1)]
        heapq.heapify(p1)
        arr = []
        for _ in range(k):
            if len(arr) == k:
                break
            if len(p1) == 0:
                break
            _, (pos1, pos2) = heapq.heappop(p1)
            arr.append([nums1[pos1], nums2[pos2]])
            if pos2+1 < len(nums2):
                heapq.heappush(
                    p1, (nums1[pos1] + nums2[pos2+1], (pos1, pos2+1)))
        return arr

Complexity

  • Time O( k * log(len(nums1)) )
Last update : October 13, 2023
Created : August 16, 2023