Program to Find K-Largest Sum Pairs in Python

When working with two lists of numbers, we often need to find the k largest sum pairs where each pair contains one element from each list. This problem can be efficiently solved using a max heap to prioritize the largest sums.

For example, if we have nums1 = [8, 6, 12], nums2 = [4, 6, 8], and k = 2, we need to find the 2 largest sum pairs. The largest pairs would be [12, 8] = 20 and [12, 6] = 18, giving us a total sum of 38.

Algorithm Approach

We'll use a max heap (implemented as min heap with negative values) to efficiently find the k largest pairs ?

• Sort both arrays in ascending order

• Start with the largest possible pair (last elements of both sorted arrays)

• Use a heap to explore pairs in descending order of their sums

• Track visited pairs to avoid duplicates

• Sum the k largest pairs

Implementation

from heapq import heappush, heappop

class Solution:
    def solve(self, nums0, nums1, k):
        if k > len(nums0) * len(nums1):
            return 0
        
        pq = []
        ans = 0
        nums0.sort()
        nums1.sort()
        i, j = len(nums0) - 1, len(nums1) - 1
        visited = set()
        
        # Start with the largest possible pair
        heappush(pq, (-(nums0[i] + nums1[j]), i, j))
        
        for _ in range(k):
            sum_val, i, j = heappop(pq)
            
            # Add adjacent pairs to explore
            if i - 1 >= 0:
                x = nums0[i - 1] + nums1[j]
                if (i - 1, j) not in visited:
                    visited.add((i - 1, j))
                    heappush(pq, (-x, i - 1, j))
            
            if j - 1 >= 0:
                y = nums0[i] + nums1[j - 1]
                if (i, j - 1) not in visited:
                    visited.add((i, j - 1))
                    heappush(pq, (-y, i, j - 1))
            
            ans += -sum_val
        
        return ans

# Test the solution
ob = Solution()
result = ob.solve([8, 6, 12], [4, 6, 8], 2)
print(f"Sum of 2 largest pairs: {result}")

# Let's also see the individual pairs
nums1 = [8, 6, 12]
nums2 = [4, 6, 8]
pairs = []
for i in nums1:
    for j in nums2:
        pairs.append((i, j, i + j))

pairs.sort(key=lambda x: x[2], reverse=True)
print(f"Top 2 pairs: {pairs[:2]}")

Sum of 2 largest pairs: 38
Top 2 pairs: [(12, 8, 20), (12, 6, 18)]

How It Works

The algorithm works by maintaining a max heap of potential pairs. Starting from the largest possible pair, it systematically explores adjacent pairs in the sorted arrays. The heap ensures we always process pairs in descending order of their sums.

Key optimizations include:

• Sorting arrays to start from maximum elements

• Using a visited set to avoid processing duplicate pairs

• Only exploring adjacent pairs to maintain efficiency

Time Complexity

The time complexity is O(k log k + n log n + m log m) where n and m are the lengths of the input arrays. The space complexity is O(k) for the heap and visited set.

Conclusion

This heap-based approach efficiently finds k largest sum pairs by exploring pairs in descending order. The algorithm avoids generating all possible pairs, making it optimal for large datasets when k is relatively small.