Program to find equal sum arrays with minimum number of operations in Python

We are given two arrays nums1 and nums2 with values between 1 and 6 (inclusive). In one operation, we can update any value in either array to any value between 1 and 6. Our goal is to find the minimum number of operations needed to make the sum of values in nums1 equal to the sum of values in nums2.

So, if the input is like nums1 = [1,5,6], nums2 = [4,1,1], then the output will be 2 because we can change nums2 from [4,1,1] to [4,1,6] in first operation, and [4,2,6] in second operation to make its sum equal to nums1.

Algorithm

To solve this, we will follow these steps ?

  • s1 := sum of all elements in nums1

  • s2 := sum of all elements in nums2

  • sort both arrays nums1 and nums2

  • if s1 > s2, then swap the arrays and their sums

  • use two pointers: left starting from 0, right starting from end of nums2

  • greedily choose the operation that gives maximum benefit in each step

Example

Let us see the implementation to get better understanding ?

def solve(nums1, nums2):
    s1 = sum(nums1)
    s2 = sum(nums2)
    nums1.sort()
    nums2.sort()
    
    if s1 > s2:
        nums1, nums2 = nums2, nums1
        s1, s2 = s2, s1

    ans = 0
    left, right = 0, len(nums2) - 1
    
    while left < len(nums1) or right >= 0:
        if s1 == s2:
            return ans
        
        curr_left = nums1[left] if left < len(nums1) else 7
        curr_right = nums2[right] if right >= 0 else 0
        
        if 6 - curr_left >= curr_right - 1:
            s1 += min(6 - curr_left, s2 - s1)
            left += 1
        else:
            s2 -= min(curr_right - 1, s2 - s1)
            right -= 1
        ans += 1
    
    return -1 if s1 != s2 else ans

nums1 = [1, 5, 6]
nums2 = [4, 1, 1]
print(solve(nums1, nums2))
2

How It Works

The algorithm uses a greedy approach:

  • We ensure s1 ? s2 by swapping if needed

  • We sort arrays to prioritize elements that can give maximum benefit

  • For nums1, we can increase small values (1,2,3) to 6 for maximum gain

  • For nums2, we can decrease large values (4,5,6) to 1 for maximum reduction

  • We choose the operation that reduces the difference most efficiently

Example Walkthrough

# Example trace
nums1 = [1, 5, 6]  # sum = 12
nums2 = [4, 1, 1]  # sum = 6

# After sorting: nums1 = [1, 5, 6], nums2 = [1, 1, 4]
# Since s1 > s2, we swap: nums1 = [1, 1, 4], nums2 = [1, 5, 6]
# Now s1 = 6, s2 = 12

print("Initial: nums1 sum =", 6, "nums2 sum =", 12)
print("Difference to overcome:", 12 - 6, "= 6")
print("Operation 1: Change 4 to 1 in nums2 ? nums2 sum = 8")
print("Operation 2: Change 1 to 4 in nums1 ? nums1 sum = 9")  
print("Final result: 2 operations needed")
Initial: nums1 sum = 6 nums2 sum = 12
Difference to overcome: 12 - 6 = 6
Operation 1: Change 4 to 1 in nums2 ? nums2 sum = 8
Operation 2: Change 1 to 4 in nums1 ? nums1 sum = 9
Final result: 2 operations needed

Conclusion

This greedy algorithm efficiently finds the minimum operations by always choosing the move that reduces the sum difference most effectively. The time complexity is O(n log n) due to sorting, where n is the total number of elements.

Updated on: 2026-03-26T14:12:50+05:30

706 Views

Kickstart Your Career

Get certified by completing the course

Get Started
Advertisements