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 instructions, which will take approximately 25 seconds. Clearly this is too long. Given the input range, we need to come up with at worst a O(nlog n) algorithm.
A better algorithm
It may be surprising that there is actually a O(n) solution. It proceeds as follows. First, we pre-compute an array is_suffix which has length n and is_suffix[i] = 1 if sum([i:]) == S/3 and 0 otherwise. This array looks like the following:
is_suffix = [0 0 0 1 1]
From this array, we compute an another array nb_suffix of length n where nb_suffix[i] = sum(is_suffix[i:]). In words, element i of nb_suffix[i] is the number of suffixes contained in arr[i:] that sum up to S/3. Concretely, nb_suffix looks like the following:
nb_suffix = [2 2 2 2 1]
Now, here's the smart bit. We begin iterating through arr and keep a running total. When we hit an index i such that the running total is equal to S/3, we do something. Since we have found the first chunk, we consider the number of partitions of the rest of the array whose sum is each S/3.
We would like to know how many partitions there is when we use arr[:i] as the first partition. nb_suffix[i+1] tells us the number of suffixes contained in arr[i+1:]. One first might think the number of suffixes using arr[:i] as the first chunk would simply be nb_suffix[i+1]. But what happens when nb_suffix[i+1] == 1 and the suffix starts at i+1? Then we would only have two chunks and would have the wrong answer. Clearly this won't work.
What about nb_suffix[i+2]? Note that even when nb_suffix[i+2] == 1 and the only suffix starts at i+2 still guarantees there is at least a middle chunk of size 1 (at arr[i+1]) whose sum must be exactly S/3. Note that this is the case because we have two other chunks whose sum is each S/3 and that when we add all parts of arr we must get back S. The only way for this to happen is if arr[i+1] == S/3.
