Program to find maximum possible value of an expression using given set of numbers in Python

We need to find the maximum possible value of an expression (sum1)² + (sum2)² where sum1 and sum2 are sums of elements from non-empty subsets of two arrays nums1 and nums2.

Given two arrays with the same number of elements, we select indices and calculate the sum of corresponding elements from each array, then find the maximum value of the sum of their squares.

Problem Understanding

For arrays nums1 = [-1, 6] and nums2 = [5, 4], possible calculations are ?

• Using index 0: (-1)² + (5)² = 1 + 25 = 26
• Using index 1: (6)² + (4)² = 36 + 16 = 52
• Using both indices: (-1 + 6)² + (5 + 4)² = 25 + 81 = 106

The maximum value is 106.

Algorithm Approach

The solution uses a geometric approach by treating pairs as 2D vectors and sorting them by their angles ?

1. Create pairs (nums1[i], nums2[i]) for each index
2. Sort pairs by their angle using atan2()
3. For each starting pair, try to extend the sum in both directions
4. Track the maximum squared magnitude found

Implementation

from math import atan2

def solve(nums1, nums2):
    # Create pairs and sort by angle
    pairs = list(zip(nums1, nums2))
    pairs = sorted(pairs, key=lambda v: atan2(v[1], v[0]))
    
    best = 0
    
    for i in range(len(pairs)):
        # Start with current pair
        u = pairs[i]
        current_sum = u[0] * u[0] + u[1] * u[1]
        
        # Try extending in forward direction
        for v in (pairs + pairs)[i + 1:(i + len(pairs))]:
            new_pair = (u[0] + v[0], u[1] + v[1])
            new_sum = new_pair[0] * new_pair[0] + new_pair[1] * new_pair[1]
            if new_sum >= current_sum:
                u = new_pair
                current_sum = new_sum
        
        if current_sum > best:
            best = current_sum
        
        # Reset and try extending in reverse direction
        u = pairs[i]
        current_sum = u[0] * u[0] + u[1] * u[1]
        
        for v in reversed((pairs + pairs)[i + 1:(i + len(pairs))]):
            new_pair = (u[0] + v[0], u[1] + v[1])
            new_sum = new_pair[0] * new_pair[0] + new_pair[1] * new_pair[1]
            if new_sum >= current_sum:
                u = new_pair
                current_sum = new_sum
                
        if current_sum > best:
            best = current_sum
    
    return best

# Test the function
nums1 = [-1, 6]
nums2 = [5, 4]
result = solve(nums1, nums2)
print(f"Maximum value: {result}")

Maximum value: 106

How It Works

The algorithm leverages the geometric property that vectors sorted by angle allow efficient exploration of combinations. By treating each (nums1[i], nums2[i]) pair as a 2D vector, we can systematically find the combination that maximizes the squared magnitude.

The double array concatenation (pairs + pairs) ensures we can wrap around and consider all possible extensions from each starting point.

Example with Different Input

# Test with different values
nums1 = [2, -3, 1]
nums2 = [4, 1, -2]
result = solve(nums1, nums2)
print(f"Maximum value: {result}")

# Show individual calculations for verification
print("\nIndividual subset calculations:")
print(f"Index 0: (2)² + (4)² = {2**2 + 4**2}")
print(f"Index 1: (-3)² + (1)² = {(-3)**2 + 1**2}")
print(f"Index 2: (1)² + (-2)² = {1**2 + (-2)**2}")
print(f"Indices 0,1: (2-3)² + (4+1)² = {(2-3)**2 + (4+1)**2}")
print(f"All indices: (2-3+1)² + (4+1-2)² = {(2-3+1)**2 + (4+1-2)**2}")

Maximum value: 26

Conclusion

This algorithm efficiently finds the maximum value by treating number pairs as 2D vectors and using geometric properties. The key insight is sorting by angles and systematically exploring extensions to find the optimal subset combination.