An explanation for my solution to 466C - Number of Ways. Credit goes to Coding Hangover for the algorithm. This post will detail the approach outlined there, explained with an example.
Problem Statement
Given an array arr, compute the number of ways we can partition arr into three chunks such that the elements in each chunk sum up up to the same number. For example the answer for arr = [1 2 3 0 3] is 2 because we can have the following splits:
[1 2][3][0 3][1 2][3 0][3]
Definitions
Define S = sum(arr).
Naive Approach
First, verify S % 3 == 0. If this is false, then it is impossible to partition arr into three equally sized chunks and the answer is 0.
Else, we continue. Pre-compute array A which is a running sum of arr:
A = [1 3 6 6 9]
Iterate i = 0...n-2 through A until we hit A[i] == S/3. At this point, we have a potential first chunk. We then iterate j = i+1...n-1 through A[i+1:] until we hit A[j] == (2/3)S. This corresponds to a potential middle chunk whose sum is S/3. Note that as long as j < n, it is guaranteed that sum(arr[j+1:]) = S/3 because we verified that sum(arr) % 3 == 0. Repeat this process for each potential first and second chunks.
Unfortunately the runtime for this approach is O(n^2) given the doubly-nested for-loop. Given that n = 5*10^5 in the worst case, we will perform (5*10^5)^2 = 25*10^10 operations, which will take approximately 250 seconds. Clearly this is too long. Given the maximum value of n, we need to cut down the runtime to O(nlogn) at worst.
A Better Algorithm
It may be surprising that there is actually a O(n) solution. The intuition is that we will still iterate i = 0...n-2 through A identifying potential first chunks. But instead of having to iterate linearly through A[i+1:] to compute the number of ways A[i+1:] can be split up into 2 chunks each having S/3 elements, we will pre-compute this value for each index and just look it up.
First, we pre-compute an array is_suffix where is_suffix[i] = 1 if sum([i:]) == S/3 and 0 otherwise:
is_suffix = [0 0 0 1 1]
Then compute an another array nb_suffix where nb_suffix[i] = sum(is_suffix[i:]). In words, nb_suffix[i] is the number of suffixes contained in arr[i:] that sum up to S/3:
nb_suffix = [2 2 2 2 1]
Now we have everything we need. We iterate i = 0..n-2 through A until we hit a potential first chunk. The number of ways to split up arr[i+1:] is simply nb_suffix[i+2]. Why i+2? Because we need a first, second, and third chunk. nb_suffix[i+2] gives us the number of last chunks. For each of these last chunks, the middle chunk is guaranteed to sum up to S/3 because we verified up front that S % 3 == 0.
For more intuition for the i+2 quantity, consider the case where nb_suffix[i+1] == 1 and is_suffix[i+1] == 1. If we used i+1 instead of i+2, then we would erroneously count this instance as a 3-patition even though the middle partition is empty.
