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.