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 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]

We will use arr as our running example throughout.

Naive Approach

First, define S = sum(arr). 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-1, it is guaranteed that sum(arr[j+1:]) = S/3 because we verified that sum(arr) % 3 == 0 up front. Hence this is a valid partitioning of arr. 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, 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(n*logn) 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 another array which will allow us to compute this value in constant time.

To this end, 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 we compute 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 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. Then the number of ways to split up arr[i+1:] into 2 chunks whose elements each sum to S/3 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 valid last chunks. For each of these last chunks (say it starts at index k), the corresponding middle chunk (which starts at j = i+1 and ends at j = k-1) is guaranteed to sum to S/3 because the first and last chunks sum to (2/3)*S and we verified up front that S % 3 == 0.