Leetcode 1712 Ways to Split Array Into Three Subarrays

Problem statement

https://leetcode.com/problems/ways-to-split-array-into-three-subarrays/

Solution

The idea is to fix left array and then use binary search to find number of options. We have two restrictions:

1. left < mid, it can be written as acc[start] >= 2*acc[i], where we need to start from index i+1.
2. mid < end, it can be written as acc[end] <= (acc[i] + acc[-1])/2, where we need to deal with last element case.

Complexity

Time complexity is O(n log n), space complexity is O(1).

Code

class Solution:
    def waysToSplit(self, nums):
        acc = list(accumulate(nums))
        ans = 0
        for i, x in enumerate(acc):
            start = bisect.bisect_left(acc, 2*x, i+1)
            end = bisect.bisect(acc, (x + acc[-1])/2)
            if end == len(nums): end -= 1
            if start < end:
                ans += (end - start)

        return ans % (10**9 + 7)