An explanation for my solution to [problem:466C]. Credit goes to [Coding Hangover](https://codinghangover.wordpress.com/2015/03/08/codeforces-266-div-2-number-of-ways/) 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`.↵
↵
For more intuition for the `i+2` quantity, consider the case where `nb_suffix[i+1] == 1` and `is_suffix[i+1] == 1`. If we used `i+1` instead of `i+2`, then we would erroneously count this instance as a 3-patition even though the middle partition is empty.
↵
#### 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`.
↵
For more intuition for the `i+2` quantity, consider the case where `nb_suffix[i+1] == 1` and `is_suffix[i+1] == 1`. If we used `i+1` instead of `i+2`, then we would erroneously count this instance as a 3-patition even though the middle partition is empty.
