split array with equal sum

Given an array with n integers, you need to find if there are triplets (i, j, k) which satisfies following conditions:

• 0 < i, i + 1 < j, j + 1 < k < n - 1
• Sum of subarrays (0, i - 1), (i + 1, j - 1), (j + 1, k - 1) and (k + 1, n - 1) should be equal.

where we define that subarray (L, R) represents a slice of the original array starting from the element indexed L to the element indexed R.

Great question, and solution (credit: lee215)

Say we are given this array:

11 2 1 2 1 2 1

We need to find these i, j, k values:

1  i   j   k
21 2 1 2 1 2 1

Because these subarrays all add up to the same value:

1  i   j   k
21 2 1 2 1 2 1
3-   -   -   -

Implementation

 1from itertools import accumulate
 2
 3class Solution:
 4    def splitArray(self, nums):
 5        acc = [*accumulate(nums)]
 6        asum = lambda i, j: acc[j] - (0, acc[i - 1])[i > 0]
 7
 8        def check(l, r):
 9            res = set()
10            for m in range(l + 1, r + 1):
11                if asum(l, m - 1) == asum(m + 1, r):
12                    res.add(asum(l, m - 1))
13            return res
14
15        for l in range(len(nums)):
16            a = check(0, l - 1)
17            b = check(l + 1, len(nums) - 1)
18            if a & b:
19                return True
20        return False
21
22
23print(Solution().splitArray([1,2,1,2,1,2,1]))

Algorithm explanation

• iterate L over nums from left to right
  • (e.g L = 3) divide the array in two parts, to the left and right of L (divide and conquer)

1        L
2[1 2 1] 2 [1 2 1]

• iterate over the array [1 2 1], if two subarrays sum to the same value (1 and 1 in this case)
  • add the result to a set, and return
• if any values in both sets match, we can subdivide this array into 4 equal parts (return True)