Intuition behind getting distance from the end of vector versus just counting

Question

I am trying to solve a GeeksForGeeks question as below:

Count the number of ways to divide an array into three contiguous parts having equal sum: Given an array of n numbers, 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. For the input {1, 2, 3, 0, 3}, the output is 2 (We have indices: (0, 1), (2, 2) and (3, 4) and (0, 1), (2, 3) and (4, 4))

Here's the code:

// 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 parts 
int 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]; 

    //<------Don't understand why we are doing things below------------->
    // 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 function 
int main() 
{ 
    int arr[] = {1 , 2 , 3 , 0 , 3}; 
    int n = sizeof(arr)/sizeof(arr[0]); 
    cout << CountContiguousParts(arr, n); 
    return 0; 
} 

I do not follow the logic behind creating a vector v storing all indexes with suffix sum equal to (total sum/3) and then doing a binary search that returns index of first element which is greater than i+1.

For the given example {1, 2, 3, 0, 3}, the prefix array is {1, 3, 6, 6, 9}, while the suffix array is {9, 8, 6, 3, 3}. So as per the above logic, for the 3 at index i=1 in the prefix array, we find 3 at index i=3 in the suffix array and add v.size() - res to count, which is our answer 2 here. However, I don't understand what this does and how it gives us our answer.

What I was thinking instead is for the the 3 at index i=1 in the prefix array, we should count the number of occurrences of 3 starting i+1 in the suffix array and return that count.

Could someone please elaborate what the above code is tryna do?

Answer 1

You want to divide an array into three parts so that the sum of each partition is equal with others. let's assume our partitions are P1, P2, P3. we want this relationship: P1=P2=P3=x

also we have P1 + P2 + P3 = 3x = total sum of array

so it is obvious that the sum of our array must be a multiple of 3 otherwise it is not possible to divide our array such that satisfies our requirements.

 vector <int> v; 
for (int i=0; i<n; i++) 
    if (suffix[i] == total_sum/3) 
        v.push_back(i);

Do you agree that P3 should contains one third of the total sum of the array? with above code we only store candidate indexes.

  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; 
        } 
    }

then with above code we also apply the same logic first we calculate where prefix includes one third of the total sum of the array lets call it i then in v (look at first code snippet what it contains) we search for an index in which its value is greater than i + 1 if we could find such an index then P2 is also possible to create lets call it j(it is the first index which is greater than i + 1) now our partitions look like this:

[0 ..... i][i+1 ....j-1][j .... n]

now what the following code do is:

    if (res != -1) 
        count += v.size() - res; 

do you agree for j and all indexes after j i.e. j+1, j+2, j+3, ... , v.size()-1 they also satisfy our requirements? i.e.

[0 ..... i] [i+1 ....v[j-1]] [v[j] .... n]

[0 ..... i] [i+1 ....v[j]]   [v[j+1] .... n]

[0 ..... i] [i+1 ....v[j + 1]] [v[j+2] .... n]
...
[0 ..... i] [i+1 ....v[v.size()-2]] [v[v.size()-1] .... n]

they all are equal partitions (v contains indexes of the original array so that if we pick the value and go to the original array and sum from there to the end its one third of total sum of the original array ). this code only counts all valid states. and the algorithm continues until it can find all states.