Grokking the Coding Interview: Patterns for Coding Questions
Ask Author
Back to course home

0% completed

Vote For New Content
Solution: Problem Challenge 1: K Pairs with Largest Sum
Table of Contents

Contents are not accessible

Contents are not accessible

Contents are not accessible

Contents are not accessible

Contents are not accessible

Problem Statement

Given two sorted arrays in descending order, find ‘K’ pairs with the largest sum where each pair consists of numbers from both the arrays.

Example 1:

Input: nums1=[9, 8, 2], nums2=[6, 3, 1], K=3
Output: [9, 3], [9, 6], [8, 6] 
Explanation: These 3 pairs have the largest sum. No other pair has a sum larger than any of these.

Example 2:

Input: nums1=[5, 2, 1], nums2=[2, -1], K=3
Output: [5, 2], [5, -1], [2, 2]

Constraints:

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

Solution

This problem follows the K-way merge pattern and we can follow a similar approach as discussed in Merge K Sorted Lists.

We can go through all the numbers of the two input arrays to create pairs and initially insert them all in the heap until we have ‘K’ pairs in Min Heap. After that, if a pair is bigger than the top (smallest) pair in the heap, we can remove the smallest pair and insert this pair in the heap.

We can optimize our algorithms in two ways:

  1. Instead of iterating over all the numbers of both arrays, we can iterate only the first ‘K’ numbers from both arrays. Since the arrays are sorted in descending order, the pairs with the maximum sum will be constituted by the first ‘K’ numbers from both the arrays.
  2. As soon as we encounter a pair with a sum that is smaller than the smallest (top) element of the heap, we don’t need to process the next elements of the array. Since the arrays are sorted in descending order, we won’t be able to find a pair with a higher sum moving forward.

Code

Here is what our algorithm will look like:

Python3
Python3
. . . .

Complexity Analysis

Time Complexity

Since, at most, we’ll be going through all the elements of both arrays and we will add/remove one element in the heap in each step, the time complexity of the above algorithm will be O(N*M*logK) where ‘N’ and ‘M’ are the total number of elements in both arrays, respectively.

If we assume that both arrays have at least ‘K’ elements then the time complexity can be simplified to O(K^2logK), because we are not iterating more than ‘K’ elements in both arrays.

Space Complexity

The space complexity will be O(K) because, at any time, our Min Heap will be storing ‘K’ largest pairs.

An Alternate Approach

The goal is to find the k pairs (nums1[i], nums2[j]) from two sorted arrays nums1 and nums2 that have the largest sums. The solution leverages a priority queue (max heap) to efficiently track the largest sums during processing.

The approach begins by pushing the sum of the first pair (nums1[0], nums2[0]) into the heap, along with its indices. The heap helps maintain the largest sums, and in each iteration, the pair with the largest sum is removed. To explore more pairs, we add new pairs formed by combining the current element in one array with the next element in the other array. Specifically, we add the pair with the next element from nums2 while keeping the current element from nums1 fixed, and vice versa. These new pairs are then pushed into the heap for further processing. This continues until we collect k pairs or the heap becomes empty.

The algorithm is efficient because the heap size is managed effectively, and only pairs relevant to finding the largest sums are processed. By iterating up to k times and leveraging the properties of the heap, the solution guarantees that the k largest sum pairs are found.

Step-by-Step Algorithm

  1. Initialize the Priority Queue:

    • Create a priority queue (maxHeap) where each entry contains:
      • sum: the sum of the two elements.
      • idx1: the index in nums1.
      • idx2: the index in nums2.

  2. Push the Initial Pair:

    • Compute the sum of the first pair (nums1[0], nums2[0]).
    • Push the pair [sum, idx1, idx2] into the heap.

  3. Process the Heap:

    • While the heap is not empty and the result size is less than k:
      • Pop the Pair with the Largest Sum:
        • Extract the top element from the heap, which represents the pair with the largest sum.
        • Add the corresponding pair (nums1[idx1], nums2[idx2]) to the result list.
      • Push Neighboring Pairs:
        • Add the pair (nums1[idx1], nums2[idx2 + 1]) if idx2 + 1 is within bounds. This step explores the next element in nums2 while keeping the current element in nums1 fixed. It ensures all possible pairs with nums1[idx1] are considered.
        • Add the pair (nums1[idx1 + 1], nums2[idx2]) if idx1 + 1 is within bounds. This step explores the next element in nums1 while keeping the current element in nums2 fixed. It ensures all possible pairs with nums2[idx2] are considered.

  4. Return the Result:

    • Once the result list contains k pairs or the heap is empty, return the list as the final output.

Algorithm Walkthrough

Input:

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

Execution:

  1. Initial Setup:

    • Priority Queue (maxHeap) = []
    • Result List (result) = []

  2. Push Initial Pair:

    • Compute sum = nums1[0] + nums2[0] = 9 + 6 = 15.
    • Push [15, 0, 0] into the heap.
    • maxHeap = [[15, 0, 0]]

  3. Iteration 1:

    • Pop Largest Pair:
      • Pop [15, 0, 0] from the heap.
      • Add [9, 6] to the result.
      • result = [[9, 6]]
    • Push Neighboring Pairs:
      • Push [12, 0, 1] for pair (9, 3).
      • Push [14, 1, 0] for pair (8, 6).
      • maxHeap = [[14, 1, 0], [12, 0, 1]]

  4. Iteration 2:

    • Pop Largest Pair:
      • Pop [14, 1, 0] from the heap.
      • Add [8, 6] to the result.
      • result = [[9, 6], [8, 6]]
    • Push Neighboring Pairs:
      • Push [11, 1, 1] for pair (8, 3).
      • Push [8, 2, 0] for pair (2, 6).
      • maxHeap = [[12, 0, 1], [11, 1, 1], [8, 2, 0]]

  5. Iteration 3:

    • Pop Largest Pair:
      • Pop [12, 0, 1] from the heap.
      • Add [9, 3] to the result.
      • result = [[9, 6], [8, 6], [9, 3]]
    • Push Neighboring Pair:
      • Push [10, 0, 2] for pair (9, 1).
      • maxHeap = [[11, 1, 1], [8, 2, 0], [10, 0, 2]]

  6. Termination:

    • The result size is now k = 3. Exit the loop.

Output:

result = [[9, 6], [8, 6], [9, 3]]

Python3
Python3
. . . .

Complexity Analysis

Time Complexity

  1. Heap Operations:

    • At most, k pairs will be added to the heap during execution.
    • Each offer (insert) and poll (remove) operation on a heap takes O(log k) time, where k is the maximum size of the heap.

  2. Iterations:

    • The while loop runs at most k times because we stop once we have found k pairs or the heap is empty.

  3. Total Complexity:

    • Each iteration involves two potential heap operations (offer), so for k iterations:
      • Total heap operations: 2 × k × O(log k).
      • Total complexity: O(k log k).

Space Complexity

  1. Heap Space:

    • At any point, the heap can contain at most k pairs.
    • Each pair in the heap is represented as an array of three integers [sum, idx1, idx2].
    • Space complexity for the heap: O(k).

  2. Result List:

    • The result list stores exactly k pairs, each represented as a List<Integer>.
    • Space complexity for the result: O(k).

  3. Auxiliary Space:

    • Constant space is used for indices and intermediate variables during computation.

  4. Total Space Complexity:

    • O(k) for the heap and result list combined.

.....

.....

.....

Like the course? Get enrolled and start learning!

Table of Contents

Contents are not accessible

Contents are not accessible

Contents are not accessible

Contents are not accessible

Contents are not accessible