I understood nlog(n) approach but I don't understand why it works. Can someone help me Understand


#1

// 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 &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];

// Storing all indexes with suffix
// sum equal to total sum by 3.
vector 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;
}