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 nums. Answer may be too large so return result modulo 10^9 + 7. Here a split of an array (with integer elements) is good if the array is split into three non-empty contiguous subarrays respectively from left to right, and the sum of the elements in left side is less than or equal to the sum of the elements in mid part, and the sum of the elements in mid part is less than or equal to the sum of the elements in right.

So, if the input is like 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]
• [2],[3,3],[3,7,1]
• [2,3],[3,3],[7,1]

To solve this, we will follow these steps −

• n := size of nums,
• m := 10^9+7
• ss := an array of size (1+n) and fill with 0
• for each index i and value val in nums , do
  • ss[i] := ss[i-1] + val
• r := 0, rr := 0, ans := 0
• for l in range 1 to n-2, do
  • r := maximum of r and l+1
  • while r < n-1 and ss[r] - ss[l] < ss[l], do
    • r := r + 1
  • rr := maximum of rr and r
  • while rr < n-1 and ss[n] - ss[rr+1] >= ss[rr+1] - ss[l], do
    • rr := rr + 1
  • if ss[l] > ss[r] - ss[l], then
    • come out from loop
  • if ss[r] - ss[l] > ss[n] - ss[r], then
    • go for next iteration
  • ans :=(ans + rr - r + 1) mod m
• return ans

Example

Let us see the following implementation to get better understanding −

def solve(nums):
   n, m = len(nums), 10**9+7
   ss = [0] * (1+n)
   for i, val in enumerate(nums, 1):
      ss[i] = ss[i-1] + val

   r = rr = ans = 0
   for l in range(1, n-1):
      r = max(r, l+1)
      while r < n-1 and ss[r]-ss[l] < ss[l]:
         r += 1
      rr = max(rr, r)
      while rr < n-1 and ss[n]-ss[rr+1] >= ss[rr+1]-ss[l]:
         rr += 1
      if ss[l] > ss[r]-ss[l]:
         break
      if ss[r]-ss[l] > ss[n]-ss[r]:
         continue
      ans = (ans+rr-r+1) % m
   return ans

nums = [2,3,3,3,7,1]
print(solve(nums))

Input

[1,7,3,6,5]

Output

3

Arnab Chakraborty

Published on 06-Oct-2021 12:39:56

• Related Questions & Answers
• Program to find number of ways to split a string in Python
• Program to find number of good ways to split a string using Python
• Program to find number of ways we can split a palindrome in python
• Program to find split a string into the max number of unique substrings in Python
• Program to find minimum number of increments on subarrays to form a target array in Python
• Program to count number of nice subarrays in Python
• Java Program to split into a number of sub-strings
• Convert array into array of subarrays - JavaScript
• Program to check a string can be split into three palindromes or not in Python
• Split number into n length array - JavaScript
• Program to find ways to make a fair array in Python
• Program to generate array by concatenating subarrays of another array in Python
• Program to find number of ways we can decode a message in Python
• Program to find sum of all odd length subarrays in Python
• Python program to split string into k distinct partitions