Balance Addicts
Contest: Codeforces Global Round 22
Difficulty rating: 2300
Tags: combinatorics, dp, math, two pointers

Given an integer sequence $$$a_1, a_2, \dots, a_n$$$ of length $$$n$$$, your task is to compute the number, modulo $$$998244353$$$, of ways to partition it into several non-empty continuous subsequences such that the sums of elements in the subsequences form a balanced sequence.

A sequence $$$s_1, s_2, \dots, s_k$$$ of length $$$k$$$ is said to be balanced, if $$$s_{i} = s_{k-i+1}$$$ for every $$$1 \leq i \leq k$$$. For example, $$$[1, 2, 3, 2, 1]$$$ and $$$[1,3,3,1]$$$ are balanced, but $$$[1,5,15]$$$ is not.

Formally, every partition can be described by a sequence of indexes $$$i_1, i_2, \dots, i_k$$$ of length $$$k$$$ with $$$1 = i_1 < i_2 < \dots < i_k \leq n$$$ such that

1. $$$k$$$ is the number of non-empty continuous subsequences in the partition;
2. For every $$$1 \leq j \leq k$$$, the $$$j$$$-th continuous subsequence starts with $$$a_{i_j}$$$, and ends exactly before $$$a_{i_{j+1}}$$$, where $$$i_{k+1} = n + 1$$$. That is, the $$$j$$$-th subsequence is $$$a_{i_j}, a_{i_j+1}, \dots, a_{i_{j+1}-1}$$$.

Let $$$s_1, s_2, \dots, s_k$$$ denote the sums of elements in the subsequences with respect to the partition $$$i_1, i_2, \dots, i_k$$$. Formally, for every $$$1 \leq j \leq k$$$, $$$$$$ s_j = \sum_{i=i_{j}}^{i_{j+1}-1} a_i = a_{i_j} + a_{i_j+1} + \dots + a_{i_{j+1}-1}. $$$$$$ For example, the partition $$$[1\,|\,2,3\,|\,4,5,6]$$$ of sequence $$$[1,2,3,4,5,6]$$$ is described by the sequence $$$[1,2,4]$$$ of indexes, and the sums of elements in the subsequences with respect to the partition is $$$[1,5,15]$$$.

Two partitions $$$i_1, i_2, \dots, i_k$$$ and $$$i'_1, i'_2, \dots, i'_{k'}$$$ (described by sequences of indexes) are considered to be different, if at least one of the following holds.

- $$$k \neq k'$$$,
- $$$i_j \neq i'_j$$$ for some $$$1 \leq j \leq \min\left\{ k, k' \right\}$$$.

Input format:
Each test contains multiple test cases. The first line contains an integer $$$t$$$ ($$$1 \leq t \leq 10^5$$$) — the number of test cases. The following lines contain the description of each test case.

The first line of each test case contains an integer $$$n$$$ ($$$1 \leq n \leq 10^5$$$), indicating the length of the sequence $$$a$$$.

The second line of each test case contains $$$n$$$ integers $$$a_1, a_2, \dots, a_n$$$ ($$$0 \leq a_i \leq 10^9$$$), indicating the elements of the sequence $$$a$$$.

It is guaranteed that the sum of $$$n$$$ over all test cases does not exceed $$$10^5$$$.

Output format:
For each test case, output the number of partitions with respect to which the sum of elements in each subsequence is balanced, modulo $$$998244353$$$.

Constraints:
Time limit: 2.0 s
Memory limit: 512.0 MB

Note:
For the first test case, there is only one way to partition a sequence of length $$$1$$$, which is itself and is, of course, balanced.

For the second test case, there are $$$2$$$ ways to partition it:

- The sequence $$$[1, 1]$$$ itself, then $$$s = [2]$$$ is balanced;
- Partition into two subsequences $$$[1\,|\,1]$$$, then $$$s = [1, 1]$$$ is balanced.

For the third test case, there are $$$3$$$ ways to partition it:

- The sequence $$$[0, 0, 1, 0]$$$ itself, then $$$s = [1]$$$ is balanced;
- $$$[0 \,|\, 0, 1 \,|\, 0]$$$, then $$$s = [0, 1, 0]$$$ is balanced;
- $$$[0, 0 \,|\, 1 \,|\, 0]$$$, then $$$s = [0, 1, 0]$$$ is balanced.

For the fourth test case, there are $$$4$$$ ways to partition it:

- The sequence $$$[1, 2, 3, 2, 1]$$$ itself, then $$$s = [9]$$$ is balanced;
- $$$[1, 2 \,|\, 3 \,|\, 2, 1]$$$, then $$$s = [3, 3, 3]$$$ is balanced;
- $$$[1 \,|\, 2, 3, 2 \,|\, 1]$$$, then $$$s = [1, 7, 1]$$$ is balanced;
- $$$[1 \,|\, 2 \,|\, 3 \,|\, 2 \,|\, 1]$$$, then $$$s = [1, 2, 3, 2, 1]$$$ is balanced.

For the fifth test case, there are $$$2$$$ ways to partition it:

- The sequence $$$[1, 3, 5, 7, 9]$$$ itself, then $$$s = [25]$$$ is balanced;
- $$$[1, 3, 5 \,|\, 7 \,|\, 9]$$$, then $$$s = [9, 7, 9]$$$ is balanced.

For the sixth test case, every possible partition should be counted. So the answer is $$$2^{32-1} \equiv 150994942 \pmod {998244353}$$$.

Example input:
6
1
1000000000
2
1 1
4
0 0 1 0
5
1 2 3 2 1
5
1 3 5 7 9
32
0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0

Example output:
1
2
3
4
2
150994942

Editorial:
Let $$$f(i, j)$$$ be the answer to the problem for the subsequence $$$a_i, a_{i+1}, \dots, a_j$$$. And we want to find $$$f(1, n)$$$.
For every $$$1 \leq i \leq j \leq n$$$, we consider the following cases.
1. $$$a_i = \dots = a_j = 0$$$. In this case, every partition produces a balanced sequence, thereby $$$f(i, j) = 2^{j-i}$$$.
2. $$$a_i = a_j = 0$$$. In this case, suppose that there are $$$x$$$ prefix $$$0$$$'s and $$$y$$$ suffix $$$0$$$'s in $$$a_i, \dots, a_j$$$ with $$$x, y \geq 1$$$. There can be at most $$$\min\{x, y\}$$$ zeros in both the prefix and suffix of the balanced sequence. There are $$$$$$ c_{x,y} = \sum_{k=0}^{\min\{x,y\}} \binom{x}{k} \binom{y}{k} $$$$$$ choices in total. So $$$f(i,j) = c_{x,y} f(i+x, j-y)$$$.
3. Otherwise, find the leftmost position $$$i \leq l \leq j$$$ and the rightmost position $$$i \leq r \leq j$$$ such that $$$a_i + \dots + a_l = a_r + \dots + a_j > 0$$$.   $$$l = j$$$. We have $$$f(i, j) = 1$$$.  $$$a_{l+1} = \dots = a_{r-1} = 0$$$. There are $$$2^{r-l}$$$ possible choices to make the sequence balanced. So $$$f(i, j) = 2^{r-l}$$$.  Otherwise, suppose there are $$$x$$$ prefix $$$0$$$'s and $$$y$$$ suffix $$$0$$$'s in $$$a_{l+1}, \dots, a_{r-1}$$$. One can cut both parts with sum $$$a_i + \dots + a_l = a_r + \dots + a_j$$$. The number of choices to cut out $$$k$$$ $$$0$$$'s in the balanced sequence is $$$$$$ \binom{x+1}{k+1} \binom{y+1}{k+1}. $$$$$$ With the case of no cut considered, we have $$$$$$ f(i, j) = \left( \sum_{k=0}^{\min\{x,y\}} \binom{x+1}{k+1} \binom{y+1}{k+1} + 1 \right) f(l+x+1, r-y-1) = c_{x+1, y+1} f(l+x+1, r-y-1). $$$$$$
4. $$$l = j$$$. We have $$$f(i, j) = 1$$$.
5. $$$a_{l+1} = \dots = a_{r-1} = 0$$$. There are $$$2^{r-l}$$$ possible choices to make the sequence balanced. So $$$f(i, j) = 2^{r-l}$$$.
6. Otherwise, suppose there are $$$x$$$ prefix $$$0$$$'s and $$$y$$$ suffix $$$0$$$'s in $$$a_{l+1}, \dots, a_{r-1}$$$. One can cut both parts with sum $$$a_i + \dots + a_l = a_r + \dots + a_j$$$. The number of choices to cut out $$$k$$$ $$$0$$$'s in the balanced sequence is $$$$$$ \binom{x+1}{k+1} \binom{y+1}{k+1}. $$$$$$ With the case of no cut considered, we have $$$$$$ f(i, j) = \left( \sum_{k=0}^{\min\{x,y\}} \binom{x+1}{k+1} \binom{y+1}{k+1} + 1 \right) f(l+x+1, r-y-1) = c_{x+1, y+1} f(l+x+1, r-y-1). $$$$$$
Since every element is processed once, it is clear that the time complexity is $$$O(n)$$$.