Check if an array of 1s and 2s can be divided into 2 parts with equal sum in Python

Suppose we have an array nums which only stores 1 and 2 in it. We have to check whether the array can be divided into two different parts such that sum of elements in each part is same.

So, if the input is like nums = [1, 1, 2, 2, 2], then the output will be True as we can divide this array like [1, 1, 2] and [2, 2] − the sum of each part is 4.

Algorithm

To solve this, we will follow these steps −

• Calculate total sum of all elements
• Count the number of 1s in the array
• If total sum is odd, return False (cannot divide into equal parts)
• If half of total sum is even, return True (can use only 2s)
• If half of total sum is odd and we have at least one 1, return True
• Otherwise, return False

Example

Let us see the following implementation to get better understanding −

def solve(nums):
    total = sum(nums)
    one_count = nums.count(1)
    
    # If total is odd, cannot divide into equal parts
    if total % 2 != 0:
        return False
    
    target = total // 2
    
    # If target sum is even, we can achieve it using only 2s
    if target % 2 == 0:
        return True
    
    # If target sum is odd, we need at least one 1
    if one_count > 0:
        return True
    
    return False

# Test the function
nums = [1, 1, 2, 2, 2]
print(solve(nums))

True

How It Works

The algorithm works by analyzing the mathematical constraints:

• Total sum check: If the total sum is odd, we cannot divide it into two equal parts
• Target sum analysis: We need each part to have sum = total/2
• Even target: If target is even, we can use only 2s to reach it
• Odd target: If target is odd, we must use at least one 1 (since 2s are even)

Additional Examples

# Test with different cases
test_cases = [
    [1, 1, 2, 2, 2],  # True: [1,1,2] and [2,2] both sum to 4
    [2, 2, 2, 2],     # True: [2,2] and [2,2] both sum to 4  
    [1, 1, 1],        # False: sum is 3, cannot divide equally
    [1, 2, 2, 2, 1]   # True: [1,2,2] and [2,1] both sum to 5
]

for i, nums in enumerate(test_cases, 1):
    result = solve(nums)
    print(f"Test {i}: {nums} ? {result}")

Test 1: [1, 1, 2, 2, 2] ? True
Test 2: [2, 2, 2, 2] ? True
Test 3: [1, 1, 1] ? False
Test 4: [1, 2, 2, 2, 1] ? True

Conclusion

This algorithm efficiently determines if an array of 1s and 2s can be split into two equal-sum parts by analyzing the total sum and count of 1s. The key insight is that odd target sums require at least one 1, while even targets can be achieved with only 2s.