Count the number of ways to divide an array into three contiguous parts having equal sum - GeeksforGeeks

19 Jan, 2022
Given an array of n numbers. Our task is to find out the number of ways to divide the array into three contiguous parts such that the sum of three parts is equal. In other words, we need to find the number of index pairs i and j such that sum of elements from 0 to i-1 is equal to the sum of elements from i to j is equal to the sum of elements from j+1 to n-1.
Examples: 
Input  : arr[] = {1, 2, 3, 0, 3}
Output : 2
Following are two possible ways to partition
1) Three parts are (1, 2), (3) and (0, 3)
2) Three parts are (1, 2), (3,0) and (3)

Input : arr[] = {0, 1, -1, 0}
Output : 1
It is only one way to partition.
1) Three parts are (0), (1,-1) and (0)
If the sum of all the elements of the array is not divisible by 3 return 0. Else it is obvious that the sum of each part of each contiguous part will be equal to the sum of all elements divided by 3.
Step 1 : Create an array of the same size as a given array whose i-th index holds the value of the sum of elements from indices 0 to i of the given array. Let’s call it prefix arrayStep 2 : Create another array of the same size as a given array whose i-th index the value of sum of elements from indices i to n-1 of the given array. Let’s call it suffix array.Step 3 : The idea is simple, we start traversing the prefix array and suppose at the i-th index of the prefix array the value of prefix array is equal to (sum of all elements of given array)/3.Step 4 : For i found in the above step we look into the suffix array from (i+2)-th index and wherever the value of suffix array is equal to (sum of all elements of given array)/3, we increase the counter variable by 1.
To implement step 4 we traverse suffix array and wherever the value of suffix array is equal to the sum of all elements of given array/3 we push that index of suffix array into the vector. And we do a binary search in the vector to calculate the number of values in suffix array which are as according to step 4.
We search in suffix array because there should be at least 1 element between the first and third part.
For more explanation see implementation below  
C++
Java
Python3
C#
Javascript
// C++ program to count number of ways we can// partition an array in three parts with equal// sum.#include<bits/stdc++.h>using namespace std; // binary search to find the number of required// indices in suffix array. Returns index of// first element which is greater than x.int binarysearch(vector <int> &v, int x){    int low = 0, high = v.size()-1;    while (low <= high)    {        int mid = (low + high)/2;        if (v[mid] <= x)            low = mid + 1;        else if (v[mid] > x && v[mid-1] <= x)            return mid;        else if (v[mid] > x && mid == 0)            return mid;        else            high = mid-1;    }    return -1;} // function to calculate the number of ways to// divide the array into three contiguous partsint CountContiguousParts(int arr[] ,int n){    int count = 0;  // initializing answer     // Creating and filling  prefix array    int prefix[n];    prefix[0] = arr[0];    for (int i=1; i<n; i++)        prefix[i] = prefix[i-1] + arr[i];     // Total sum of elements is equal to last    // value in prefix array.    int total_sum = prefix[n-1];     // If sum of all the elements is not divisible    // by 3, we can't divide array in 3 parts of    // same sum.    if (total_sum%3 != 0)        return 0;     // Creating and filling suffix array    int suffix[n];    suffix[n-1] = arr[n-1];    for (int i=n-2; i>=0; i--)        suffix[i] = suffix[i+1] + arr[i];     // Storing all indexes with suffix    // sum equal to total sum by 3.    vector <int> v;    for (int i=0; i<n; i++)        if (suffix[i] == total_sum/3)            v.push_back(i);     // Traversing the prefix array and    // doing binary search in above vector    for (int i=0; i<n; i++)    {        // We found a prefix with total_sum/3        if (prefix[i] == total_sum/3)        {            // Find first index in v[] which            // is greater than i+1.            int res = binarysearch(v, i+1);             if (res != -1)                count += v.size() - res;        }    }     return count;} // driver functionint main(){    int arr[] = {1 , 2 , 3 , 0 , 3};    int n = sizeof(arr)/sizeof(arr[0]);    cout << CountContiguousParts(arr, n);    return 0;}
// Java program to count number of ways we can// partition an array in three parts with equal// sum.import java.util.*; class GFG{ // binary search to find the number of required// indices in suffix array. Returns index of// first element which is greater than x.static int binarysearch(Vector<Integer> v, int x){    int low = 0, high = v.size() - 1;    while (low <= high)    {        int mid = (low + high) / 2;        if (v.get(mid) <= x)            low = mid + 1;        else if (v.get(mid) > x &&                 v.get(mid) <= x)            return mid;        else if (v.get(mid) > x && mid == 0)            return mid;        else            high = mid - 1;    }    return -1;} // function to calculate the number of ways to// divide the array into three contiguous partsstatic int CountContiguousParts(int arr[], int n){    int count = 0; // initializing answer     // Creating and filling prefix array    int []prefix = new int[n];    prefix[0] = arr[0];    for (int i = 1; i < n; i++)        prefix[i] = prefix[i - 1] + arr[i];     // Total sum of elements is equal to last    // value in prefix array.    int total_sum = prefix[n - 1];     // If sum of all the elements is not divisible    // by 3, we can't divide array in 3 parts of    // same sum.    if (total_sum % 3 != 0)        return 0;     // Creating and filling suffix array    int []suffix = new int[n];    suffix[n - 1] = arr[n - 1];    for (int i = n - 2; i >= 0; i--)        suffix[i] = suffix[i + 1] + arr[i];     // Storing all indexes with suffix    // sum equal to total sum by 3.    Vector<Integer> v = new Vector<>();    for (int i = 0; i < n; i++)        if (suffix[i] == total_sum / 3)            v.add(i);     // Traversing the prefix array and    // doing binary search in above vector    for (int i = 0; i < n; i++)    {        // We found a prefix with total_sum/3        if (prefix[i] == total_sum / 3)        {            // Find first index in v[] which            // is greater than i+1.            int res = binarysearch(v, i + 1);             if (res != -1)                count += v.size() - res;        }    }    return count;} // Driver Codepublic static void main(String[] args){    int arr[] = {1 , 2 , 3 , 0 , 3};    int n = arr.length;    System.out.println(CountContiguousParts(arr, n));}} // This code is contributed by Princi Singh
# Python 3 program to count the number of ways we can# partition an array in three parts with equal# sum. # binary search to find the number of required# indices in suffix array. Returns index of# first element which is greater than x.def binarysearch(v, x):    low = 0    high = len(v) - 1    while (low <= high):        mid = int((low + high) / 2)        if (v[mid] <= x):            low = mid + 1        elif (v[mid] > x and v[mid - 1] <= x):            return mid        elif (v[mid] > x and mid == 0):            return mid        else:            high = mid - 1     return -1 # function to calculate the number of ways to# divide the array into three contiguous partsdef CountContiguousParts(arr,n):    count = 0 # initializing answer     # Creating and filling prefix array    prefix = [0 for i in range(n)]    prefix[0] = arr[0]    for i in range(1, n, 1):        prefix[i] = prefix[i - 1] + arr[i]     # Total sum of elements is equal to last    # value in prefix array.    total_sum = prefix[n - 1]     # If sum of all the elements is not divisible    # by 3, we can't divide array in 3 parts of    # same sum.    if (total_sum % 3 != 0):        return 0     # Creating and filling suffix array    suffix = [0 for i in range(n)]    suffix[n - 1] = arr[n - 1]    i = n - 2    while(i >= 0):        suffix[i] = suffix[i + 1] + arr[i]        i -= 1     # Storing all indexes with suffix    # sum equal to total sum by 3.    v = []    for i in range(n):        if (suffix[i] == int(total_sum / 3)):            v.append(i)     # Traversing the prefix array and    # doing binary search in above vector    for i in range(n):                 # We found a prefix with total_sum/3        if (prefix[i] == int(total_sum / 3)):                         # Find first index in v[] which            # is greater than i+1.            res = binarysearch(v, i + 1)             if (res != -1):                count += len(v) - res     return count # Driver Codeif __name__ == '__main__':    arr = [1 , 2 , 3 , 0 , 3]    n = len(arr)    print(CountContiguousParts(arr, n)) # This code is contributed by# Surendra_Gangwar
// C# program to count number of ways we can// partition an array in three parts with equalusing System;using System.Collections.Generic; class GFG{ // binary search to find the number of required// indices in suffix array. Returns index of// first element which is greater than x.static int binarysearch(List<int> v, int x){    int low = 0, high = v.Count - 1;    while (low <= high)    {        int mid = (low + high) / 2;        if (v[mid] <= x)            low = mid + 1;        else if (v[mid] > x &&                v[mid] <= x)            return mid;        else if (v[mid] > x && mid == 0)            return mid;        else            high = mid - 1;    }    return -1;} // function to calculate the number of ways to// divide the array into three contiguous partsstatic int CountContiguousParts(int []arr, int n){    int count = 0; // initializing answer     // Creating and filling prefix array    int []prefix = new int[n];    prefix[0] = arr[0];    for (int i = 1; i < n; i++)        prefix[i] = prefix[i - 1] + arr[i];     // Total sum of elements is equal to last    // value in prefix array.    int total_sum = prefix[n - 1];     // If sum of all the elements is not divisible    // by 3, we can't divide array in 3 parts of    // same sum.    if (total_sum % 3 != 0)        return 0;     // Creating and filling suffix array    int []suffix = new int[n];    suffix[n - 1] = arr[n - 1];    for (int i = n - 2; i >= 0; i--)        suffix[i] = suffix[i + 1] + arr[i];     // Storing all indexes with suffix    // sum equal to total sum by 3.    List<int> v = new List<int>();    for (int i = 0; i < n; i++)        if (suffix[i] == total_sum / 3)            v.Add(i);     // Traversing the prefix array and    // doing binary search in above vector    for (int i = 0; i < n; i++)    {        // We found a prefix with total_sum/3        if (prefix[i] == total_sum / 3)        {            // Find first index in v[] which            // is greater than i+1.            int res = binarysearch(v, i + 1);             if (res != -1)                count += v.Count - res;        }    }    return count;} // Driver Codepublic static void Main(String[] args){    int []arr = {1 , 2 , 3 , 0 , 3};    int n = arr.Length;    Console.WriteLine(CountContiguousParts(arr, n));}} // This code is contributed by Rajput-Ji
<script> // Javascript program to count number// of ways we can partition an array// in three parts with equal sum.     // Binary search to find the number of required// indices in suffix array. Returns index of// first element which is greater than x.function binarysearch(v, x){    let low = 0, high = v.length - 1;         while (low <= high)    {        let mid = Math.floor((low + high) / 2);                 if (v[mid] <= x)            low = mid + 1;        else if (v[mid] > x && v[mid] <= x)            return mid;        else if (v[mid] > x && mid == 0)            return mid;        else            high = mid - 1;    }    return -1;} // Function to calculate the number of ways to// divide the array into three contiguous partsfunction CountContiguousParts(arr, n){         // Initializing answer    let count = 0;         // Creating and filling prefix array    let prefix = new Array(n);    prefix[0] = arr[0];         for(let i = 1; i < n; i++)        prefix[i] = prefix[i - 1] + arr[i];         // Total sum of elements is equal to last    // value in prefix array.    let total_sum = prefix[n - 1];         // If sum of all the elements is not divisible    // by 3, we can't divide array in 3 parts of    // same sum.    if (total_sum % 3 != 0)        return 0;         // Creating and filling suffix array    let suffix = new Array(n);    suffix[n - 1] = arr[n - 1];         for(let i = n - 2; i >= 0; i--)        suffix[i] = suffix[i + 1] + arr[i];         // Storing all indexes with suffix    // sum equal to total sum by 3.    let v = [];    for(let i = 0; i < n; i++)        if (suffix[i] == Math.floor(total_sum / 3))            v.push(i);         // Traversing the prefix array and    // doing binary search in above vector    for(let i = 0; i < n; i++)    {                 // We found a prefix with total_sum/3        if (prefix[i] == Math.floor(total_sum / 3))        {                         // Find first index in v[] which            // is greater than i+1.            let res = binarysearch(v, i + 1);                         if (res != -1)                count += v.length - res;        }    }    return count;} // Driver Codelet arr = [ 1 , 2 , 3 , 0 , 3 ];let  n = arr.length; document.write(CountContiguousParts(arr, n)); // This code is contributed by avanitrachhadiya2155 </script>
Output: 
2
Time Complexity is O(n log n)
Efficient Approach [ O(n) solution ] :
If sum of all the elements of the array is not divisible by 3, return 0.It is obvious that the sum of each part of each contiguous part will be equal to the sum of all elements divided by 3.Let’s create an array cnt[ ], where cnt[ i ] equals 1, if the sum of elements from i-th to n-th equals Array_Sum/3 else 0. Now, calculate the cumulative sum of the cnt array from the last index.Thus, we receive very simple solution: for each prefix of initial array 1...i with the sum that equals Array_Sum/3 we need to add to the answer sums[ i+2 ].
If sum of all the elements of the array is not divisible by 3, return 0.
It is obvious that the sum of each part of each contiguous part will be equal to the sum of all elements divided by 3.
Let’s create an array cnt[ ], where cnt[ i ] equals 1, if the sum of elements from i-th to n-th equals Array_Sum/3 else 0. Now, calculate the cumulative sum of the cnt array from the last index.
Thus, we receive very simple solution: for each prefix of initial array 1...i with the sum that equals Array_Sum/3 we need to add to the answer sums[ i+2 ].
Below is the code for the above approach.  
C
Java
Python3
C#
Javascript
// C++ program to count the ways// to break the array in 3 equal parts// having equal sum.#include <bits/stdc++.h>using namespace std; // Function to count the no of waysint countways(int a[], int n){    int cnt[n] = {0}, s = 0;         // Calculating the sum of the array    // and storing it in variable s    for (int i = 0 ; i < n ; i++)    {        s += a[i];    }     // Checking s is divisible by 3 or not    if (s % 3 != 0)        return 0;         // Calculating the sum of each part    s /= 3;         int ss = 0;         // If the sum of elements from i-th to n-th    // equals to sum of part putting 1 in cnt    // array otherwise 0.    for (int i = n-1; i >= 0 ; i--)    {        ss += a[i];        if (ss == s)            cnt[i] = 1;    }         // Calculating the cumulative sum    // of the array cnt from the last index.    for (int i = n-2 ; i >= 0 ; i--)        cnt[i] += cnt[i + 1];         int ans = 0;    ss = 0;         // Calculating answer using original    // and cnt array.    for (int i = 0 ; i+2 < n ; i++)    {        ss += a[i];        if (ss == s)            ans += cnt[i + 2];    }    return ans;} // Driver functionint main(){    int n = 5;    int a[] = {1, 2, 3, 0, 3};    cout << countways(a, n) << endl;    return 0;}
// Java program to count the ways to// break the array in 3 equal parts// having equal sum.import java.io.*; class GFG {         // Function to count the no of ways    static int countways(int a[], int n)    {        int cnt[] = new int[n];        int s = 0;                 // Calculating the sum of the array        // and storing it in variable s        for (int i = 0 ; i < n ; i++)        {            s += a[i];        }             // Checking s is divisible by 3 or not        if (s % 3 != 0)            return 0;                 // Calculating the sum of each part        s /= 3;                 int ss = 0;                 // If the sum of elements from i-th to n-th        // equals to sum of part then storing the number          // of times possible in the cnt array.        for (int i = n-1; i >= 0 ; i--)        {            ss += a[i];            if (ss == s)            {                if(i == n-1)                {                  cnt[i] = 1;                }                  else                {                  cnt[i] = 1 + cnt[i+1];                }            }        }                 int ans = 0;        ss = 0;                 // Calculating answer using original        // and cnt array.        for (int i = 0 ; i+2 < n ; i++)        {            ss += a[i];            if (ss == s)                ans += cnt[i + 2];        }        return ans;    }         // Driver function    public static void main (String[] args)    {        int n = 5;        int a[] = {1, 2, 3, 0, 3};        System.out.println(countways(a, n));    }} // This code is contributed by anuj_67.
# Python3 program to count the ways# to break the array in 3 equal parts# having equal sum. # Function to count the no of waysdef countways(a, n):     cnt = [0 for i in range(n)]    s = 0     # Calculating the sum of the array    # and storing it in variable s    s = sum(a)     # Checking s is divisible by 3 or not    if (s % 3 != 0):        return 0     # Calculating the sum of each part    s //= 3     ss = 0     # If the sum of elements from i-th    # to n-th equals to sum of part    # putting 1 in cnt array otherwise 0.    for i in range(n - 1, -1, -1):         ss += a[i]        if (ss == s):            cnt[i] = 1     # Calculating the cumulative sum    # of the array cnt from the last index.    for i in range(n - 2, -1, -1):        cnt[i] += cnt[i + 1]     ans = 0    ss = 0     # Calculating answer using original    # and cnt array.    for i in range(0, n - 2):        ss += a[i]        if (ss == s):            ans += cnt[i + 2]     return ans # Driver Coden = 5a = [1, 2, 3, 0, 3]print(countways(a, n)) # This code is contributed# by mohit kumar
// C# program to count the ways to// break the array in 3 equal parts// having an equal sum.using System; class GFG{         // Function to count the no of ways    static int countways(int[] a, int n)    {        int[] cnt = new int[n];        int s = 0;                 // Calculating the sum of the array        // and storing it in variable s        for (int i = 0 ; i < n ; i++)        {            s += a[i];        }             // Checking s is divisible by 3 or not        if (s % 3 != 0)            return 0;                 // Calculating the sum of each part        s /= 3;                 int ss = 0;                 // If the sum of elements from i-th to n-th        // equals to sum of part putting 1 in cnt        // array otherwise 0.        for (int i = n - 1; i >= 0 ; i--)        {            ss += a[i];            if (ss == s)                cnt[i] = 1;        }                 // Calculating the cumulative sum        // of the array cnt from the last index.        for (int i = n - 2 ; i >= 0 ; i--)            cnt[i] += cnt[i + 1];                 int ans = 0;        ss = 0;                 // Calculating answer using original        // and cnt array.        for (int i = 0 ; i + 2 < n ; i++)        {            ss += a[i];            if (ss == s)                ans += cnt[i + 2];        }        return ans;    }         // Driver code    public static void Main ()    {        int n = 5;        int[] a = {1, 2, 3, 0, 3};        Console.Write(countways(a, n));    }} // This code is contributed by Ita_c.
<script> // Javascript program to count the ways to// break the array in 3 equal parts// having equal sum.         // Function to count the no of ways    function countways(a,n)    {        let cnt = new Array(n);        let s = 0;                  // Calculating the sum of the array        // and storing it in variable s        for (let i = 0 ; i < n ; i++)        {            s += a[i];        }              // Checking s is divisible by 3 or not        if (s % 3 != 0)            return 0;                  // Calculating the sum of each part        s = Math.floor(s/3);                  let ss = 0;                  // If the sum of elements from i-th to n-th        // equals to sum of part putting 1 in cnt        // array otherwise 0.        for (let i = n-1; i >= 0 ; i--)        {            ss += a[i];            if (ss == s)                cnt[i] = 1;        }                  // Calculating the cumulative sum        // of the array cnt from the last index.        for (let i = n-2 ; i >= 0 ; i--)            cnt[i] += cnt[i + 1];                  let ans = 0;        ss = 0;                  // Calculating answer using original        // and cnt array.        for (let i = 0 ; i+2 < n ; i++)        {            ss += a[i];            if (ss == s)                ans += cnt[i + 2];        }        return ans;    }         // Driver function    let n = 5;    let a=[1, 2, 3, 0, 3];    document.write(countways(a, n));           // This code is contributed by rag2127     </script>
Output: 
2
Time Complexity: O(n). 
This approach is contributed by Abhishek Sharma. This article is contributed by Ayush Jha. If you like GeeksforGeeks and would like to contribute, you can also write an article using write.geeksforgeeks.org or mail your article to review-team@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.Please write comments if you find anything incorrect, or you want to share more information about the topic discussed above. 
vt_m
ukasp
mohit kumar 29
SURENDRA_GANGWAR
princi singh
Rajput-Ji
geekyamitesh
rag2127
avanitrachhadiya2155
RajSuvariya
Binary Search
prefix-sum
Arrays
prefix-sum
Arrays
Binary Search
Writing code in comment?
Please use ide.geeksforgeeks.org,
generate link and share the link here.
Maximum and minimum of an array using minimum number of comparisons
Top 50 Array Coding Problems for Interviews
Stack Data Structure (Introduction and Program)
Introduction to Arrays
Multidimensional Arrays in Java
Linear Search
Linked List vs Array
Given an array A[] and a number x, check for pair in A[] with sum as x (aka Two Sum)
Python | Using 2D arrays/lists the right way
Search an element in a sorted and rotated array