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:

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:

1 1 2 1 2 1 2 1

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

1   i   j   k
2 1 2 1 2 1 2 1

Because these subarrays all add up to the same value:

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

Implementation

 1 from itertools import accumulate
 2 
 3 class 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 
23 print(Solution().splitArray([1,2,1,2,1,2,1]))

Algorithm explanation

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