This is documentation for my solution to 466C - Number of Ways. Credit goes to Coding Hangover for the algorithm. This post will detail the approach here, as well as provide an example.
Problem Statement
Given an array arr
, we are to compute the number of ways we can partition arr
into 3 chunks where each chunk is contiguous. Further, the sum of the elements in each chunk must be the same.
Running Example
We will use the input array arr = [1 2 3 0 3]
as a running example.
Naive Approach
My first approach to this problem was to pre-compute an array A
which is a running sum of arr
. That is, A[i] = sum(a[:i])
inclusive. Concretely, we would have the following array:
A = [1 3 6 6 9]
.
We would then iterate through A
. Define S = sum(arr)
. Every time we reach index i
such that A[i] == S/3
, we would iterate from j = i+1
until we reach index j
such that A[j] == (2/3)S
. This corresponds to a "middle" chunk which sums up to S/3
. To verify the last chunk sums up to S/3
, we would compute A[n] - A[j]
. If this number equals S/3
, then we have found a valid 3-chunk.
Unfortunately the runtime for this approach is O(n^2)
. 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
.