Given an array that has n integers, count the number of ways to split all elements of this array into 3 contiguous parts so that the sum of each part is the same. Each part must not be empty.
Algorithm: O(N) runtime
1. If total sum % 3 != 0, return 0.
2. Keep a count of prefix sums that are 1 / 3 of the total sum, denote it by psOneThird.
3. Visit all elements from index 1 to n - 2 and for each newly visited element a[i], do the following:
(a). If the current sum is 2 / 3 of total sum, we know we just found a partition that leaves the unvisited elements summing up to 1 / 3 of total sum. Since we just find this partition by adding a[i] to the current sum, we can conclude that there are psOneThird distinct pairs that have 1 / 3 of total sum each, with the right half of 1 / 3 ends right after a[i]. Add psOneThird to the final answer.
(b). If the current sum is 1 / 3 of total sum, update psOneThird.
Why does the above algorithm work?
Each valid partition must have the 2nd part ends on an element from a[1] to a[n - 2]. And each time the 2nd part ends on one of these elements, it generates some more unique partitions. (The 2nd part is different). So we check from a[1] to a[n - 2]. After adding one element to the current sum and it becomes 2 /3 of the total sum, we've just found such a unique partition. Now we just need to add the count of (1/3, 2/3) pairs, contributed by this newly found 2nd part. Since the 2nd part is fixed, the count is just how many prefix sums of 1 / 3 of total sum we've seen so far.
One key note here is that we must check if the current sum is 2/3 of total sum first, then check if 1/3. This ensures we don't have 1/3 and 2/3 partition ends on the same spot, right after a[i], for the corner case of 0 total sum.
private static void solve(int q, FastScanner in, PrintWriter out) {
int n = in.nextInt();
int[] a = new int[n];
long sum = 0;
for(int i = 0; i < n; i++) {
a[i] = in.nextInt();
sum += a[i];
}
if(sum % 3 != 0) {
out.println(0);
}
else {
long currSum = a[0], oneThirdSum = sum / 3, ans = 0;
int psOneThird = (a[0] == oneThirdSum ? 1 : 0);
for(int i = 1; i < n - 1; i++) {
currSum += a[i];
if(currSum == oneThirdSum * 2) {
ans += psOneThird;
}
if(currSum == oneThirdSum) {
psOneThird++;
}
}
out.println(ans);
}
out.close();
}