Hints: Counting Prefixes
Forget about the array $a$ for a moment and look at the prefix sums.
The sequence of prefix sums of $a$ is
$$0 = p_0,\ p_1,\ p_2,\ \ldots,\ p_n,$$where consecutive values differ by exactly $\pm 1$ (since each $a_i \in \{-1, +1\}$). So the unsorted prefix-sum sequence is exactly a walk on the integer line:
- it starts at $0$,
- it makes $n$ steps, each of size $\pm 1$.
The input gives you the multiset of values that this walk visits at times $1, 2, \ldots, n$ (the value at time $0$ is fixed to $0$).
Question. Reformulate the problem purely as a walk-counting problem. What “type” of object is the input describing about the walk?
Answer to Hint 1. The input is exactly the multiset $\{p_1, p_2, \ldots, p_n\}$ — i.e. the visit counts of the walk at every integer value, excluding the value at time $0$ (which is always $0$).
Define $\mathrm{cnt}[v] = $ number of indices $i \in \{0, 1, \ldots, n\}$ at which the walk is at value $v$. Then
$$\mathrm{cnt}[v] = (\text{number of }p_i\text{ equal to }v) + [v = 0].$$So the problem becomes:
Count walks $v_0, v_1, \ldots, v_n$ on $\mathbb{Z}$ with $v_0 = 0$, $|v_{i+1} - v_i| = 1$, such that for every $v \in \mathbb{Z}$, the value $v$ appears exactly $\mathrm{cnt}[v]$ times in $v_0, v_1, \ldots, v_n$.
Two trivial sanity checks first:
- The set $S = \{v : \mathrm{cnt}[v] > 0\}$ must form a contiguous interval of integers (the walk on $\mathbb{Z}$ moves by $\pm 1$, so the visited values are an interval).
- That interval must contain $0$.
If either fails, the answer is $0$.
Question. The walk has a known start ($0$) but the endpoint $v_n$ is not given. How can we deal with that?
Answer to Hint 2. Iterate over the endpoint!
Let $f$ be the candidate value of $v_n$ (the endpoint of the walk, equivalently the final prefix sum, equivalently the sum of the array $a$). For the answer to be non-zero we need $\mathrm{cnt}[f] \ge 1$.
Then
$$\text{answer} = \sum_{f\,:\,\mathrm{cnt}[f] \ge 1}\ N(f),$$where $N(f)$ is the number of walks with start $0$, end $f$, and visit multiset $\mathrm{cnt}$.
So we have reduced the problem to the easy version:
Easy version. Given visit counts $\mathrm{cnt}[v]$ on a contiguous integer interval, a fixed start $s$ and a fixed end $f$ (both inside the interval), count the number of $\pm 1$ walks consistent with all of this.
The full answer just sums the easy version over all valid endpoints $f$. If we can solve the easy version in $O(n)$, the whole problem is $O(n^2)$, which is fast enough for $n \le 5000$.
Question (easy version). With both endpoints fixed and visit counts fixed, how many “up” steps ($v \to v+1$) and “down” steps ($v+1 \to v$) does the walk make for each adjacent pair $(v, v+1)$?
Answer to Hint 3. Define for each $i$ in the visited interval:
- $\mathrm{go}_R[i]$ = number of “up” steps from $i$ to $i+1$,
- $\mathrm{go}_L[i]$ = number of “down” steps from $i+1$ to $i$.
These are the edge multiplicities on a multigraph whose vertices are the visited integers and whose only edges connect consecutive integers.
For each vertex $v$, the number of times the walk enters $v$ equals the number of times it leaves $v$, with the only exceptions being the start ($+1$ extra leaving) and the end ($+1$ extra entering). Concretely:
$$\underbrace{\mathrm{go}_R[v-1] + \mathrm{go}_L[v]}_{\text{enters }v} = \mathrm{cnt}[v] - [v = s],$$ $$\underbrace{\mathrm{go}_L[v-1] + \mathrm{go}_R[v]}_{\text{leaves }v} = \mathrm{cnt}[v] - [v = f].$$(Treat $\mathrm{go}_R[v-1]$ and $\mathrm{go}_L[v-1]$ as $0$ when $v$ is the leftmost visited integer.)
Sweeping $v$ from left to right we can compute $\mathrm{go}_R[v]$ and $\mathrm{go}_L[v]$ uniquely from the equations above:
$$\mathrm{go}_L[v] = \bigl(\mathrm{cnt}[v] - [v = s]\bigr) - \mathrm{go}_R[v-1],$$ $$\mathrm{go}_R[v] = \bigl(\mathrm{cnt}[v] - [v = f]\bigr) - \mathrm{go}_L[v-1].$$So the edge counts are completely forced by $(s, f, \mathrm{cnt})$. No freedom there.
Question. What are the validity conditions on the resulting $\mathrm{go}_L, \mathrm{go}_R$?
Answer to Hint 4. Two requirements:
- Non-negativity: $\mathrm{go}_L[v] \ge 0$ and $\mathrm{go}_R[v] \ge 0$ for all $v$.
- Right boundary: at the rightmost visited integer $r$, both $\mathrm{go}_L[r]$ and $\mathrm{go}_R[r]$ must be $0$ (no edges past the rightmost visited vertex).
If any of these fails, $N(f) = 0$.
The right-boundary condition can equivalently be stated as: when we finish the left-to-right sweep, the recurrence must “close cleanly” — no leftover edges sticking out.
Question. Suppose the edge counts $\mathrm{go}_L, \mathrm{go}_R$ are non-negative and consistent. Now we have a directed multigraph on a line: vertex $v$ has $\mathrm{go}_R[v]$ up-edges to $v+1$ and $\mathrm{go}_L[v]$ down-edges to $v$. We want to count the number of distinct walks from $s$ to $f$ that traverse every edge exactly once (i.e. Eulerian trails). How do we count those?
Answer to Hint 5. This is a classical Eulerian trail counting problem on a multigraph.
For our specific graph (vertices on a line, edges only between consecutive integers), there is a clean way to count.
Think of the walk as a sequence of decisions:
- Each time the walk arrives at vertex $v$, it must leave again — except the final visit at $v = f$.
- When leaving $v$, it picks one of $\mathrm{go}_L[v-1] + \mathrm{go}_R[v]$ outgoing edges (some go left, some go right).
So at vertex $v$ the walk makes
$$o_v := \mathrm{cnt}[v] - [v = f]$$exit decisions. Among those, $\mathrm{go}_R[v]$ must be “go right” and $\mathrm{go}_L[v-1]$ must be “go left”. Two walks are different iff their sequence of left/right decisions at some vertex differs.
If we just naively multiply
$$\prod_v \binom{o_v}{\mathrm{go}_R[v]} = \prod_v \frac{o_v!}{\mathrm{go}_R[v]!\,\mathrm{go}_L[v-1]!}$$we slightly over-count: not every assignment of left/right decisions corresponds to a complete Eulerian trail (some can get “stuck”). The correct adjustment is given by the BEST theorem (or, equivalently, by counting in-arborescences rooted at the end).
Question (BEST theorem on a line graph). In a path-graph multigraph, what is the number of in-arborescences rooted at vertex $f$?
Answer to Hint 6. An in-arborescence rooted at $f$ is a choice, for every non-root vertex $v$, of one outgoing edge such that following these edges always eventually reaches $f$.
On a line graph this is trivial:
- For $v < f$, the chosen outgoing edge must go right (otherwise you’d diverge from $f$). So you pick one of $\mathrm{go}_R[v]$ up-edges.
- For $v > f$, the chosen outgoing edge must go left. So you pick one of $\mathrm{go}_L[v-1]$ down-edges.
Hence
$$T_f \;=\; \prod_{v < f}\mathrm{go}_R[v]\ \cdot\ \prod_{v > f}\mathrm{go}_L[v-1].$$The BEST theorem says: the number of Eulerian trails from $s$ to $f$ equals
$$N(f) \;=\; T_{\!f}\ \cdot\ \prod_{v}\bigl(\mathrm{outdeg}(v) - 1\bigr)!\ \cdot\ \frac{\mathrm{outdeg}(s)}{\prod_{v}\bigl(\text{ordering inside outgoing-edge classes}\bigr)},$$which simplifies on a path graph to a very tidy product. The reusable form is:
$$\boxed{\;N(f) \;=\; T_{\!f}\ \cdot\ \prod_{v}\frac{(\mathrm{cnt}[v] - 1)!}{\mathrm{go}_R[v]!\,\mathrm{go}_L[v]!}\;}$$where the factorial-product is the BEST-theorem multinomial term (treating “left” and “right” outgoing edges at each vertex as indistinguishable within their class), and $T_f$ from above is the arborescence factor.
(This matches the formula in model_sol.cpp — see Hint 9 for the exact mapping. If you have not seen the BEST theorem, you can also derive the same formula by induction on the number of edges, processing the rightmost vertex first; or you can derive it via the more elementary “stars and bars” recurrence sketched in the official editorial.)
Question. What’s a useful sanity check / alternative derivation, in case the BEST theorem feels heavy?
Answer to Hint 7 (alternative derivation). Here is an elementary reformulation that gives the same count, level by level. (This is the version in maomao90’s official editorial.)
For a fixed sum $f$, start with the array
$$a^{(0)} = \underbrace{[1, 1, \ldots, 1]}_{p_n\text{ ones}},\ \underbrace{[-1, -1, \ldots, -1]}_{p_n - f\text{ minus ones}}.$$Its prefix sums climb $1, 2, \ldots, p_n$ and then descend $p_n - 1, p_n - 2, \ldots, f$. So initially:
- value $p_n$ is hit once,
- each value strictly between $f$ and $p_n$ is hit twice,
- value $f$ is hit once.
Now, inserting a pair $(-1, +1)$ into the array immediately after a position whose current prefix sum is $v$ creates two new prefix sums: one of value $v - 1$ and one of value $v$. Everything else is unchanged.
Process the levels from top down: $v = p_n, p_n - 1, \ldots$
- Suppose at level $v$ we currently have $x$ occurrences of value $v$ but the input requires $y$ (with $y \ge x$).
- We need to pick $y - x$ insertion slots, each immediately after one of the $x$ existing positions where the prefix sum is $v$ (with repetition allowed). By stars and bars, the number of ways is
- Each such insertion also increases the count at level $v - 1$ by $1$. So when we move down to the next level, we update $x_{v-1} \leftarrow x_{v-1}^{(\text{original})} + (y - x)$.
If at any level $x > y$, this $f$ contributes $0$. Otherwise multiply the per-level binomials. This is also $O(n)$ per fixed $f$.
Answer to Hint 8 (mapping the BEST formula to the code). Let’s tour model_sol.cpp:
-
int l = min(p[0], 0); int r = max(p[n - 1], 0);— bound the visited interval. If $r - \ell > n$ the walk cannot span the interval in $n$ steps, output $0$. -
cnt[p[i] - l] += 1;thencnt[s] += 1;— build $\mathrm{cnt}$ on shifted coordinates, where $s = -\ell$ is the shifted start (i.e. value $0$). The extra+1at $s$ accounts for the implicit $p_0 = 0$ visit. -
The outer loop
for (int f = 0; f <= r - l; f++)iterates the candidate endpoint (in shifted coordinates). -
The inner sweep computes $\mathrm{go}_L[i], \mathrm{go}_R[i]$ exactly by the recurrence in Hint 4, and sets
failif a count goes negative or if the right boundary doesn’t close to $\mathrm{go}_L = \mathrm{go}_R = 0$. -
The two product loops
for (int i = s; i > 0; i--) cur *= go_l[i - 1]; for (int i = s; i < r - l; i++) cur *= go_r[i];compute one form of the arborescence factor in the BEST theorem applied to the Eulerian augmentation $G + (f \to s)$. On a line graph this is just one chosen multiplicity per adjacent-vertex pair; all valid choices of root give the same count, so the picked form is convenient: take a “down” multiplicity for edges left of $s$ and an “up” multiplicity for edges right of (and at) $s$.
-
The final two loops
for (int i = 0; i <= r - l; i++) cur *= fact[cnt[i] - 1]; for (int i = 0; i < r - l; i++) { cur *= inv_fact[go_l[i]]; cur *= inv_fact[go_r[i]]; }realize the multinomial factor $\prod_v (\mathrm{cnt}[v] - 1)! / (\mathrm{go}_R[v]!\,\mathrm{go}_L[v]!)$, which counts the orderings of “left vs right” exits at each vertex.
Adding cur to ans over all valid $f$ gives the final answer modulo $998\,244\,353$. (Sanity check the formula by computing ans by hand on the sample $p = [-1, 0, 0, 1, 1]$ and confirming you get $3$.)
Putting everything together:
- Read $p_1 \le \ldots \le p_n$.
- Let $\ell = \min(p_1, 0)$, $r = \max(p_n, 0)$. If $r - \ell > n$, output $0$ (the walk can’t reach that wide).
- Build $\mathrm{cnt}[v]$ on the shifted range $[0, r - \ell]$, including the extra visit at the starting value $0$. If the support is not contiguous (there is a gap with $\mathrm{cnt}[v] = 0$ inside the interval), output $0$.
- Set $s = -\ell$ (the shifted start). For each candidate end $f \in [0, r - \ell]$ with $\mathrm{cnt}[f] > 0$:
- Sweep left to right computing $\mathrm{go}_L[v], \mathrm{go}_R[v]$ from the conservation equations of Hint 4. Bail out if any value is negative or if the right boundary doesn’t close to $0$.
- Compute $N(f)$ via the formula of Hint 7 (or equivalently of Hint 9).
- Add $N(f)$ to the answer.
- Print the answer modulo $998\,244\,353$.
Complexity. Each fixed $f$ takes $O(n)$ time for the sweep and the product. There are $O(n)$ candidate endpoints, so the total work per test case is $O(n^2)$. With $\sum n \le 5000$, the worst case is well under $2.5 \cdot 10^7$ operations and runs comfortably in the time limit.
Key takeaways.
- Reformulate. “Sorted prefix sums of a $\pm 1$ array” is exactly “visit multiset of a $\pm 1$ walk on $\mathbb Z$”.
- Pin down a free parameter. Fix the endpoint of the walk and sum.
- Edges before walks. For fixed endpoints and visit counts, the edge multiplicities of the walk are forced; the only remaining freedom is the order in which exits at each vertex are chosen, and that freedom is counted by the BEST theorem (equivalently a stars-and-bars layer-by-layer construction).
- Precompute factorials and inverse factorials so each $N(f)$ collapses to a clean $O(n)$ product.