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`

.

Auto comment: topic has been updated by ebanner (previous revision, new revision, compare).i couldn't get it from i+2 point in the last 3rd line that says..... "Why i+2? Because we need ".......please be a bit discriptive

Good explanation, thanks mate!

Dude, it already has an editorial (that kinda says the same thing)

How to solve this problem with DP?pls