1712. Ways to Split Array Into Three Subarrays - cocoder39/coco39_LC GitHub Wiki

1712. Ways to Split Array Into Three Subarrays

notes:

  1. straightforward to come up with prefix sum + binary search
  2. trick to consider boundary condition
    • 2.1 left [0, len(nums)-3]
    • 2.2 left + 1 <= midLowerBound < midUpperBound <= len(nums)-1
      • 2.2.1 midLowerBound = max(midLowerBound, left+1)
      • 2.2.2 midUpperBound = min(midUpperBound, len(nums)-1)
      • 2.2.3 count += max(0, midUpperBound - midLowerBound)
class Solution:
    def waysToSplit(self, nums: List[int]) -> int:
        prefixSum = 0
        prefixSums = []
        for num in nums:
            prefixSum += num
            prefixSums.append(prefixSum)
        
        count = 0
        MOD = 10**9 + 7
        for left in range(len(nums)-2): # [0, len(nums)-3]
            # prefixSums[left] <= prefixSums[midLowerBound] - prefixSums[left] 
            midLowerBound = bisect.bisect_left(prefixSums, 2 * prefixSums[left]) 
            # prefixSums[midUpperBound] - prefixSums[left] <= prefixSums[-1] - prefixSums[midUpperBound] 
            midUpperBound = bisect.bisect_right(prefixSums, (prefixSums[left] + prefixSums[-1])//2)
            # left + 1 <= midLowerBound < midUpperBound <= len(nums)-1
            midLowerBound = max(midLowerBound, left+1)
            midUpperBound = min(midUpperBound, len(nums)-1)
            count += max(0, midUpperBound - midLowerBound)
            count %= MOD
        return count