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
.