Home/

...

/

Problem Challenge 1

Problem Challenge 1

Dry-run templates

This is the implementation of the solution provided for the problem posed in the Problem Challenge 1 lesson:

Press + to interact
Python 3.5
from __future__ import print_function
from heapq import *
def find_k_largest_pairs(nums1, nums2, k):
  minHeap = []
  for i in range(0, min(k, len(nums1))):
    for j in range(min(k, len(nums2))):
      if len(minHeap) < k:
        heappush(minHeap, (nums1[i] + nums2[j], i, j))
      else:
        # if the sum of the two numbers from the two arrays is smaller than the smallest(top)
        # element of the heap, we can 'break' here. Since the arrays are sorted in the
        # descending order, we'll not be able to find a pair with a higher sum moving forward
        if nums1[i] + nums2[j] < minHeap[0][0]:
          break
        else:  # we have a pair with a larger sum, remove top and insert this pair in the heap
          heappop(minHeap)
          heappush(minHeap, (nums1[i] + nums2[j], i, j))
  result = []
  for (num, i, j) in minHeap:
    result.append([nums1[i], nums2[j]])
  return result
def main():
  print("Pairs with largest sum are: " +
        str(find_k_largest_pairs([9, 8, 2], [6, 3, 1], 3)))
main()

Students may be encouraged to run through the provided solution with the following sample inputs:

Sample Input #1

nums1: [9, 8, 2]
nums2: [6, 3, 1]
k: 3

Iteration No.

Line No.

minHeap

( nums1[i] , nums2[j] )

minHeap[0][0]

result

-

6

[]

-

-

-

1

9-10

[(15, 0, 0)]

(9,6)

15

-

2

9-10

[(12, 0, 1), (15, 0, 0)]

(9,3)

12

-

3

9-10

[(10, 0, 2), (15, 0, 0), (12, 0, 1)]

(9,1)

10

-

4

17-19

[(12, 0, 1), (15, 0, 0), (14, 1, 0)]

(8,6)

12

-

5

15-16

[(12, 0, 1), (15, 0, 0), (14, 1, 0)]

(8,3)

12

-

6

15-16

[(12, 0, 1), (15, 0, 0), (14, 1, 0)]

(2,6)

12

-

1

33-34

[(12, 0, 1), (15, 0, 0), (14, 1, 0)]

(9,3)

-

[[9, 3]]

2

33-34

[(12, 0, 1), (15, 0, 0), (14, 1, 0)]

(9,6)

-

[[9, 3], [9, 6]]

3

33-34

[(12, 0, 1), (15, 0, 0), (14, 1, 0)]

(8,6)

-

[[9, 3], [9, 6], [8, 6]]

Sample Input #2

nums1: [3, 3, 3]
nums2: [3, 3, 3]
k: 2

Iteration No.

Line No.

minHeap

( nums1[i] , nums2[j] )

minHeap[0][0]

result

-

6

[]

-

-

-

1

9-10

[(6, 0, 0)]

(3,3)

(6, 0, 0)

-

2

9-10

[(6, 0, 0), (6, 0, 1)]

(3,3)

(6, 0, 0)

-

3

17-19

[(6, 0, 1), (6, 1, 0)]

(3,3)

(6, 0, 1)

-

4

17-19

[(6, 1, 0), (6, 1, 1)]

(3,3)

(6, 1, 0)


1

33-34

[(6, 1, 0), (6, 1, 1)]

(3,3)

-

[[3, 3]]

2

33-34

[(6, 1, 0), (6, 1, 1)]

(3,3)

-

[[3, 3], [3, 3]]

Step-by-step solution construction

Let’s first write the code to push all pairs in minHeap.

Lines 8-13 iterate over the two arrays and push each pair ( nums1[i], nums2[j] ) on to minHeap.

Press + to interact
Python 3.5
from __future__ import print_function
from heapq import *
def find_k_largest_pairs(nums1, nums2, k):
  minHeap = []
  iteration=0
  for i in range(0, min(k, len(nums1))):
    for j in range(min(k, len(nums2))):
      print("Iteration: "+str(iteration+1))
      iteration +=1
      print("\tPush next pair: " + "(" + str(nums1[i]) + ", " + str(nums2[j]) + ")")
      heappush(minHeap, (nums1[i] + nums2[j], i, j))
      print("\tminHeap after adding pair is: "+str(minHeap))
  result = []
  return result
def main():
  arr = [[[9, 8, 2], [6, 3, 1],3],[[3, 3, 3], [3, 3, 3],2]]
  for i in range(len(arr)):
    print("nums1 arrray: " + str(arr[i][0]))
    print("nums2 array: " + str(arr[i][1]))
    print("k: "+str(arr[i][2]) + "\n")
    result = find_k_largest_pairs(arr[i][0], arr[i][1], arr[i][2])
    print("\npairs with largest sum are: " + str(result))
    print("-"*100)
main()

Now, let’s modify the code to insert only the first k pairs in the minHeap and filter the rest of the pairs.

Lines 12-14 push the next pair on to minHeap only if its size is less than k. Lines 17-21 collect the top k pairs from minHeap.

Press + to interact
Python 3.5
from __future__ import print_function
from heapq import *
def find_k_largest_pairs(nums1, nums2, k):
  minHeap = []
  iteration=0
  for i in range(0, min(k, len(nums1))):
    for j in range(min(k, len(nums2))):
      print("Iteration: "+str(iteration+1))
      iteration +=1
      if len(minHeap) < k:
        print("\tNot yet made " + str(k) + " pairs, so push next pair: " + "(" + str(nums1[i]) + ", " + str(nums2[j]) + ")")
        heappush(minHeap, (nums1[i] + nums2[j], i, j))
        print("\tminHeap after adding pair is: "+str(minHeap))
  result = []
  for (num, i, j) in minHeap:
    print("Adding pair " + "(" + str(nums1[i]) + "," + str(nums2[j]) + ") to the results array")
    result.append([nums1[i], nums2[j]])
    print("result: " + str(result))
  return result
def main():
  arr = [[[9, 8, 2], [6, 3, 1],3],[[3, 3, 3], [3, 3, 3],2]]
  for i in range(len(arr)):
    print("nums1 arrray: " + str(arr[i][0]))
    print("nums2 array: " + str(arr[i][1]))
    print("k: "+str(arr[i][2]) + "\n")
    result = find_k_largest_pairs(arr[i][0], arr[i][1], arr[i][2])
    print("\npairs with largest sum are: " + str(result))
    print("-"*100)
main()

However, this is not the correct result, as, in the first test case, the pair (8,6)(8, 6) has not been considered. This is happening because we are currently only pushing the first k pairs on to the heap. There may well be pairs after this point whose sum is greater. To safely consider those pairs, we extend our logic in the code below, checking each subsequent pair for the following two conditions:

  • Lines 25-28: If the sum of the new pair is larger than, or equal to, the lowest sum in the minHeap, push it on to the heap. Before pushing, perform a pop on the minHeap so that the pair with the lowest sum is removed.

  • Lines 20-24: If the sum of the new pair is lower than the lowest sum in the minHeap, break out of the loop, as the arrays are sorted in descending order and we will not be able to find a pair with a higher sum moving forward beyond this point.

Press + to interact
Python 3.5
from __future__ import print_function
from heapq import *
 
 
def find_k_largest_pairs(nums1, nums2, k):
 minHeap = []
 iteration=0
 for i in range(0, min(k, len(nums1))):
   for j in range(min(k, len(nums2))):
     print("Iteration: "+str(iteration+1))
     iteration +=1
     if len(minHeap) < k:
       print("\tNot yet made " + str(k) + " pairs, so push next pair: " + "(" + str(nums1[i]) + ", " + str(nums2[j]) + ")")
       heappush(minHeap, (nums1[i] + nums2[j], i, j))
       print("\tminHeap after adding pair is: "+str(minHeap))
     else:
       # if the sum of the two numbers from the two arrays is smaller than the smallest (top)
       # element of the heap, we can 'break' here. Since the arrays are sorted in
       # descending order, we will not be able to find a pair with a higher sum moving forward
       if nums1[i] + nums2[j] < minHeap[0][0]:
         print("\tImpossible to find pairs with greater sums, breaking out of the loop")
         print("\tminHeap at this point: "+str(minHeap))
         print("\tBreaking at pair: " + "(" + str(nums1[i]) + ", " + str(nums2[j]) + ")")
         break
       else:  # we have a pair with a larger sum, remove top and insert this pair in the heap
         print("\tPair with a larger sum is FOUND, remove top and push this pair on to the heap")
         heappop(minHeap)
         heappush(minHeap, (nums1[i] + nums2[j], i, j))
         print("\tPush new pair: " + "(" + str(nums1[i]) + ", " + str(nums2[j]) + ")")
         print("\tminHeap after pop and push: "+str(minHeap))
 
 result = []
 for (num, i, j) in minHeap:
   print("Adding pair " + "(" + str(nums1[i]) + ", " + str(nums2[j]) + ") to the results array")
   result.append([nums1[i], nums2[j]])
   print("result: " + str(result))
 
 return result
 
 
 
def main():
 arr = [[[9, 8, 2], [6, 3, 1],3],[[3, 3, 3], [3, 3, 3],2]]
 for i in range(len(arr)):
   print("nums1 arrray: " + str(arr[i][0]))
   print("nums2 array: " + str(arr[i][1]))
   print("k: "+str(arr[i][2]) + "\n")
   result = find_k_largest_pairs(arr[i][0], arr[i][1], arr[i][2])
   print("\nPairs with the largest sums are: " + str(result))
   print("-"*100)
 
main()