Program to find number of ways to split array into three subarrays in Python

Suppose we have an array called nums, we have to find the number of good ways to split this array. A split is considered good if the array is divided into three non-empty contiguous subarrays where sum(left) ? sum(middle) ? sum(right). Since the answer may be large, return result modulo 10^9 + 7.

Problem Understanding

If the input is nums = [2,3,3,3,7,1], then the output will be 3 because there are three different ways of splitting ?

• [2],[3],[3,3,7,1] ? sums: 2 ? 3 ? 11
• [2],[3,3],[3,7,1] ? sums: 2 ? 6 ? 11
• [2,3],[3,3],[7,1] ? sums: 5 ? 6 ? 8

Algorithm

To solve this, we will follow these steps ?

• Create a prefix sum array to calculate subarray sums efficiently
• For each possible left boundary, find valid middle and right boundaries
• Use two pointers to efficiently count valid splits
• Return the count modulo 10^9 + 7

Solution

def solve(nums):
    n, m = len(nums), 10**9 + 7
    
    # Create prefix sum array
    prefix_sum = [0] * (n + 1)
    for i, val in enumerate(nums):
        prefix_sum[i + 1] = prefix_sum[i] + val
    
    r = rr = ans = 0
    
    # Try each possible left subarray ending position
    for l in range(1, n - 1):
        # Find minimum valid middle subarray start
        r = max(r, l + 1)
        while r < n - 1 and prefix_sum[r] - prefix_sum[l] < prefix_sum[l]:
            r += 1
        
        # Find maximum valid middle subarray end
        rr = max(rr, r)
        while rr < n - 1 and prefix_sum[n] - prefix_sum[rr + 1] >= prefix_sum[rr + 1] - prefix_sum[l]:
            rr += 1
        
        # Check if left sum is too large
        if prefix_sum[l] > prefix_sum[r] - prefix_sum[l]:
            break
        
        # Check if middle sum is too large
        if prefix_sum[r] - prefix_sum[l] > prefix_sum[n] - prefix_sum[r]:
            continue
        
        # Add count of valid middle subarrays
        ans = (ans + rr - r + 1) % m
    
    return ans

# Test the function
nums = [2, 3, 3, 3, 7, 1]
print(f"Input: {nums}")
print(f"Output: {solve(nums)}")

# Test with another example
nums2 = [1, 7, 3, 6, 5]
print(f"Input: {nums2}")
print(f"Output: {solve(nums2)}")

Input: [2, 3, 3, 3, 7, 1]
Output: 3
Input: [1, 7, 3, 6, 5]
Output: 2

How It Works

The algorithm uses a prefix sum array to efficiently calculate subarray sums. For each possible left subarray ending position l, it finds the range of valid middle subarray positions using two pointers r and rr. The key insight is that once we fix the left boundary, we can efficiently find all valid middle boundaries that satisfy both conditions.

Time Complexity

The time complexity is O(n) where n is the length of the array, as each pointer moves at most n times throughout the entire algorithm. The space complexity is O(n) for the prefix sum array.

Conclusion

This solution efficiently finds all valid three-way splits using prefix sums and two-pointers technique. The algorithm ensures that each split satisfies the condition sum(left) ? sum(middle) ? sum(right) while maintaining optimal time complexity.