Hamed and AghaBalaSar

Let $M=a_n$. Since the array is snake, $M$ is the global maximum, and every earlier value is in $[0,M]$.

The nasty-looking part is the pseudocode. The counting is not the first boss fight; understanding what the loop actually sums is.

What does the loop compute?

Look at the indices visited by the algorithm:

$1=i_1<i_2<\dots<i_k=n.$

Some transitions are jumps to nxt[pos], and those jumps strictly increase the value. Some transitions are ordinary pos += 1, which happen only when the current value already equals $M$.

Split the visited sequence right after every visited position whose value is $M$. Inside each block, the jump values strictly increase, so that block contributes:

$\text{last value of block}-\text{first value of block}.$

Adding the blocks telescopes into:

$$f(a)=\sum \text{values at visited global maxima} -\sum \text{values immediately after those maxima} -a_1.$$

The last maximum is $a_n$, so there is no “value immediately after it”. This formula is the key. The loop is basically charging every climb, not doing anything mystical. Classic 2800 problem behavior: the code is scarier than the math, because of course it is.

Now we sum this formula over all snake arrays.

Bounded compositions

Define

$g(k,s,x)$

as the number of arrays of length $k$ with sum $s$ and all elements between $0$ and $x$.

Without the upper bound, the answer is $\binom{s+k-1}{k-1}$. To enforce $a_i\le x$, use inclusion-exclusion on elements that exceed $x$:

$$g(k,s,x)=\sum_{t=0}^{\lfloor s/(x+1)\rfloor} (-1)^t\binom{k}{t}\binom{s+k-1-t(x+1)}{k-1}.$$

If $k=0$, then $g(0,s,x)=1$ only when $s=0$, otherwise $0$.

This can be computed in $O(s/(x+1))$ time.

Fix the maximum

Fix $a_n=x$. Then the first $n-1$ elements have sum $m-x$ and each is at most $x$. So the number of such arrays is

$G_1=g(n-1,m-x,x).$

We also need the count where a chosen earlier position is forced to equal $x$. After fixing that position and $a_n=x$, the remaining $n-2$ positions have sum $m-2x$, still bounded by $x$, so define

$G_2=g(n-2,m-2x,x).$

Now compute the three parts of the formula.

1. Sum of visited global maxima

Every valid array contributes the last element $a_n=x$, giving

$xG_1.$

Also, any earlier position with value $x$ is a global maximum that the algorithm visits by ordinary movement through max positions. For each of the $n-1$ earlier positions, forcing it to be $x$ gives $G_2$ arrays. Therefore this part contributes

$xG_1+x(n-1)G_2.$

2. Subtract $a_1$

Among the first $n-1$ positions, everything is symmetric. Across all arrays counted by $G_1$, the total sum of those positions is always $m-x$. Therefore each position contributes the same total sum:

$\frac{m-x}{n-1}G_1.$

So the total contribution of $a_1$ is

$\frac{m-x}{n-1}G_1.$

Modulo $10^9+7$, division by $n-1$ means multiplying by its modular inverse.

3. Subtract values immediately after earlier global maxima

Take an index $i$ with $2\le i<n$. If $a_{i-1}=x$, then $a_i$ is immediately after a global maximum. Fixing $a_{i-1}=x$ and $a_n=x$, the remaining sum is $m-2x$ over $n-2$ positions. By symmetry, the total sum over all choices of these middle following positions is simply

$(m-2x)G_2.$

For $i=n$, if $a_{n-1}=x$, the value immediately after that maximum is $a_n=x$, contributing

$xG_2.$

So the whole “after maximum” subtraction is

$xG_2+(m-2x)G_2.$

This formula still works when $n=2$: the middle part is empty, and $g(0,s,x)$ handles it.

Putting it together

For every $x$ from $0$ to $m$:

$$G_1=g(n-1,m-x,x),\qquad G_2=g(n-2,m-2x,x).$$

Add

$$xG_1+x(n-1)G_2$$

and subtract

$$\frac{m-x}{n-1}G_1+xG_2+(m-2x)G_2.$$

All terms are modulo $10^9+7$.

The inclusion-exclusion loop for a fixed $x$ takes $O(m/(x+1))$, and we call it twice, so the total complexity is

$$O\left(\sum_{x=0}^{m}\frac{m}{x+1}\right)=O(m\log m).$$

Since the sum of $m$ over tests is at most $2\cdot 10^5$, this fits easily.