educative.io

Solution with O(K*LogK) time complexity

I take K elements from the first array and pair them with the first element of the second array and put them in a max_heap. Then I get a pair for the heap and advance in the second array according to that particular pair.
The worst case is when the result is the first element of the first array paired with K distinct elements from the second array, but that would be 2KLogK, better than K^2LogK.

def find_k_largest_pairs(nums1, nums2, k):
  max_heap = []
  for i in range(min(k, len(nums1))):
    heappush(max_heap, (-nums1[i] - nums2[0], i, 0))
  results = []
  while len(results) < k and max_heap:
    _, i, j = heappop(max_heap)
    results.append([nums1[i], nums2[j]])
    j += 1
    if j < len(nums2):
      heappush(max_heap, (-nums1[i] - nums2[j], i, j))
  return results

Any thoughts?

Type your question above this line.

Course: https://www.educative.io/collection/5668639101419520/5671464854355968
Lesson: https://www.educative.io/collection/page/5668639101419520/5671464854355968/6521584850305024

Hello @Andres_Parra,

Thanks for reaching out to us with a better solution. Yes, you are right that the above solution’s time complexity is O(2KLogK), which is better than the provided solution’s time complexity O(K^2LogK). We will update our solution after this suggestion.

Thank You. :blush:

1 Like