No. of ways to divide an array

Question

I want to find The number of ways to divide an array into 3 contiguous parts such that the sum of the three parts is equal

-10^9 <= A[i] <= 10^9

My approach: Taking Input and Checking for Base Case:

for(int i=0;i<n;i++){
     a[i]= in.nextLong();
     sum+=a[i];
}
if(sum%3!=0)System.out.println("0");

If The answer is not above Then Forming the Prefix and Suffix Sum.

for(int i=1;i<=n-2;i++){
        xx+=a[i-1];
        if(xx==sum/3){

            dp[i]=1;
        }        
    }

Suffix Sum and Updating the Binary Index Tree:

for(int i=n ;i>=3;i--){
         xx+=a[i-1];
         if(xx==sum/3){
             update(i, 1, suffix);
         }
    }

And Now simple Looping the array to find the Total Ways: int ans=0;

for(int i=1;i<=n-2;i++){

        if(dp[i]==1)
        {

            ans+=  (query(n, suffix) - query(i+1, suffix)); 
         // Checking For the Sum/3 in array where index>i+1
        }
}

I Getting the wrong answer for the above approach

I don't Know where I have made mistake Please Help to Correct my mistake.

Update and Query Function:

public static void update(int i , int value , int[] arr){

       while(i<arr.length){
           arr[i]+=value;

           i+=i&-i;

       }
}
public static int query(int i ,int[] arr){

     int ans=0;

       while(i>0){

           ans+=arr[i];
           i-=i&-i;
       }

       return ans;

}

Answer 1

As far as your approach is concerned its correct. But there are some points because of which it might give WA

1. Its very likely that sum overflows int as each element can magnitude of 10^9, so use long long .
2. Make sure that suffix and dp array are initialized to 0.

Having said that using a BIT tree here is an overkill , because it can be done in O(n) compared to your O(nlogn) solution ( but does not matter if incase you are submitting on a online judge ). For the O(n) approach just take your suffix[] array.And as you have done mark suffix[i]=1 if sum from i to n is sum/3, traversing the array backwards this can be done in O(n). Then just traverse again from backwards doing suffix[i]+=suffix[i-1]( apart from base case i=n).So now suffix[i] stores number of indexs i<=j<=n such that sum from index j to n is sum/3, which is what you are trying to achieve using BIT.
So what I suggest either write a bruteforce or this simple O(n) and check your code against it, because as far as your approach is concerned it is correct, and debugging is something not suited for stackoverflow.

Answer 2

First, we calculate an array dp, with dp[i] = sum from 0 to i, this can be done in O(n)

long[]dp = new long[n];
for(int i = 0; i < n; i++)
   dp[i] = a[i];
   if(i > 0)
     dp[i] += dp[i - 1];

Second, let say the total sum of array is x, so we need to find at which position, we have dp[i] == x/3;

For each i position which have dp[i] == 2*x/3, we need to add to final result, the number of index j < i, which dp[j] == x/3.

int count = 0;
int result = 0;
for(int i = 0; i < n - 1; i++){
    if(dp[i] == x/3)
       count++;
    else if(dp[i] == x*2/3)
       result += count;
}

The answer is in result.

What wrong with your approach is,

    if(dp[i]==1)
    {

        ans+=  (query(n, suffix) - query(i+1, suffix)); 
     // Checking For the Sum/3 in array where index>i+1
    }

This is wrong, it should be

(query(n, suffix) - query(i, suffix));

Because, we only need to remove those from 1 to i, not 1 to i + 1.

Not only that, this part:

for(int i=1;i<=n-2;i++){
  //....      
}

Should be i <= n - 1;

Similarly, this part, for(int i=n ;i>=3;i--), should be i >= 1

And first part:

for(int i=0;i<n;i++){
     a[i]= in.nextLong();
     sum+=a[i];
}

Should be

for(int i=1;i<=n;i++){
     a[i]= in.nextLong();
     sum+=a[i];
}

A lot of small errors in your code, which you need to put in a lot of effort to debugging first, jumping to ask here is not a good idea.