Tokitsukaze and Gems
Contest: Codeforces Round 789 (Div. 1)
Difficulty rating: 3500
Tags: dp, math

Tokitsukaze has a sequence with length of $$$n$$$. She likes gems very much. There are $$$n$$$ kinds of gems. The gems of the $$$i$$$-th kind are on the $$$i$$$-th position, and there are $$$a_i$$$ gems of the same kind on that position. Define $$$G(l,r)$$$ as the multiset containing all gems on the segment $$$[l,r]$$$ (inclusive).

A multiset of gems can be represented as $$$S=[s_1,s_2,\ldots,s_n]$$$, which is a non-negative integer sequence of length $$$n$$$ and means that $$$S$$$ contains $$$s_i$$$ gems of the $$$i$$$-th kind in the multiset. A multiset $$$T=[t_1,t_2,\ldots,t_n]$$$ is a multisubset of $$$S=[s_1,s_2,\ldots,s_n]$$$ if and only if $$$t_i\le s_i$$$ for any $$$i$$$ satisfying $$$1\le i\le n$$$.

Now, given two positive integers $$$k$$$ and $$$p$$$, you need to calculate the result of

$$$$$$\sum_{l=1}^n \sum_{r=l}^n\sum\limits_{[t_1,t_2,\cdots,t_n] \subseteq G(l,r)}\left(\left(\sum_{i=1}^n p^{t_i}t_i^k\right)\left(\sum_{i=1}^n[t_i>0]\right)\right),$$$$$$

where $$$[t_i>0]=1$$$ if $$$t_i>0$$$ and $$$[t_i>0]=0$$$ if $$$t_i=0$$$.

Since the answer can be quite large, print it modulo $$$998\,244\,353$$$.

Input format:
The first line contains three integers $$$n$$$, $$$k$$$ and $$$p$$$ ($$$1\le n \le 10^5$$$; $$$1\le k\le 10^5$$$; $$$2\le p\le 998\,244\,351$$$) — the length of the sequence, the numbers $$$k$$$ and $$$p$$$.

The second line contains $$$n$$$ integers $$$a_1, a_2, \ldots, a_n$$$ ($$$1\le a_i\le 998\,244\,351$$$) — the number of gems on the $$$i$$$-th position.

Output format:
Print a single integers — the result modulo $$$998\,244\,353$$$.

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

Example input:
5 2 2
1 1 1 2 2

Example output:
6428

Example input:
6 2 2
2 2 2 2 2 3

Example output:
338940

Editorial:
First, consider how to calculate the contribution of a legal path. It can be found that the contribution of a legal path is equivalent to the following formula, where $$$b_i$$$ indicates the number of the $$$i$$$-th gems on the path.
$$$$$$ \sum_{i_1=0}^{b_1}\sum_{i_2=0}^{b_2}\ldots \sum_{i_n=0}^{b_n}(\sum_{j=1}^n (P^{i_j}i_j^k)\times\sum_{j=1}^n [i_j\ne 0]) $$$$$$
This formula seems impossible to calculate, but it can be observed that there is a chain structure behind it. Consider dynamic programming, try to calculate the $$$1\sim (i+1)$$$-th gems' answer by only considering the $$$1\sim i$$$-th gems' answer. Consider maintenance these:
$$$dp_{n,1}=\sum_{i_1=0}^{b_1}\sum_{i_2=0}^{b_2}\ldots \sum_{i_n=0}^{b_n}(\sum_{j=1}^n (P^{i_j}i_j^k)\times\sum_{j=1}^n [i_j\ne 0])$$$
$$$dp_{n,2}=\sum_{i_1=0}^{b_1}\sum_{i_2=0}^{b_2}\ldots \sum_{i_n=0}^{b_n}\sum_{j=1}^n (P^{i_j}i_j^k)$$$
$$$dp_{n,3}=\sum_{i_1=0}^{b_1}\sum_{i_2=0}^{b_2}\ldots \sum_{i_n=0}^{b_n} \sum_{j=1}^n [i_j\ne 0]$$$
$$$dp_{n,4}=\sum_{i_1=0}^{b_1}\sum_{i_2=0}^{b_2}\ldots \sum_{i_n=0}^{b_n} 1$$$
It can be observed that if we enumerate how many the $$$(n + 1)$$$-th gem are selected, this can actually be transferred, for example:
$$$dp_{n+1,1}=\sum_{p=0}^{b_{n+1}}\sum_{i_1=0}^{b_1}\sum_{i_2=0}^{b_2}\ldots \sum_{i_n=0}^{b_n}((\sum_{j=1}^n P^{i_j}i_j^k+P^p p^k)\times (\sum_{j=1}^n [i_j\ne 0])+[p\ne 0])$$$
Then spread.
$$$A=\sum_{p=0}^{b_{n+1}}\sum_{i_1=0}^{b_1}\sum_{i_2=0}^{b_2}\ldots \sum_{i_n=0}^{b_n}(\sum_{j=1}^n P^{i_j}i_j^k \times \sum_{j=1}^n [i_j\ne 0])=(1+b_{n+1})dp_{n,1}$$$
$$$B=\sum_{p=0}^{b_{n+1}}\sum_{i_1=0}^{b_1}\sum_{i_2=0}^{b_2}\ldots \sum_{i_n=0}^{b_n}\sum_{j=1}^n P^{i_j}i_j^k[p\ne 0]=b_{n+1}dp_{n,2}$$$
$$$C=\sum_{p=0}^{b_{n+1}}\sum_{i_1=0}^{b_1}\sum_{i_2=0}^{b_2}\ldots \sum_{i_n=0}^{b_n}P^p p^k\times \sum_{j=1}^n [i_j\ne 0]=dp_{n,3} \times \sum_{p=0}^{b_{n+1}}P^p p^k=dp_{n,3} \times \sum_{p=1}^{b_{n+1}}P^p p^k$$$
$$$D=\sum_{p=0}^{b_{n+1}}\sum_{i_1=0}^{b_1}\sum_{i_2=0}^{b_2}\ldots \sum_{i_n=0}^{b_n}(P^p p^k)\times [p\ne 0]=dp_{n,4}\times\sum_{p=1}^{b_{n+1}}P^p p^k$$$
Now it can be transferred. The main problem is how to calculate $$$\sum_{i=1}^n i^kx^i$$$ quickly, ensuring that $$$x\gt 1$$$.
Let $$$h_{k,n}$$$ represent $$$\sum_{i=1}^n i^k x^i$$$.
Calculate the sum according to the proportional sequence, $$$h_{0,n}=\dfrac{x^{n+1}-x}{x-1}$$$.
$$$$$$ xh_{k,n}-h_{k,n}=\sum_{i=1}^n x^{i+1}i^k-\sum_{i=1}^nx^ii^k $$$$$$
$$$$$$ =x^{n+1}n^k+\sum_{i=1}^nx^i((i-1)^k-i^k) $$$$$$
$$$$$$ =x^{n+1}n^k+\sum_{i=1}^nx^i([\sum_{j=0}^k(-1)^{k-j}\binom{k}{j}i^j]-i^k) $$$$$$
$$$$$$ =x^{n+1}n^k+\sum_{i=0}^{k-1}(-1)^{k-j}\binom{k}{j}h_{i,n} $$$$$$
Finally we get $$$$$$ h_{k,n}=\dfrac{x^{n+1}n^k+\sum_{i=0}^{k-1}(-1)^{k-j}\binom{k}{j}h_{i,n}}{x-1} $$$$$$ According to this recurrence formula, we can get $$$$$$ h_{0,n}=x^{n+1}(\frac{1}{x-1})-x\times(\frac{1}{x-1}) $$$$$$
$$$$$$ h_{1,n}=x^{n+1}(\frac{n}{x-1}-\dfrac{1}{(x-1)^2})-x(-\dfrac{1}{(x-1)^2}) $$$$$$
When $$$k$$$ is constant, there exist a $$$k$$$-degree polynomial $$$g_k(x)$$$ about $$$n$$$, making $$$h_{k,n}=x^{n+1}g_k(n)-xg_k(0)$$$. This conclusion can be proved by mathematical induction.
Consider how to use interpolation to get this polynomial, we only need to find $$$g_k(0)\sim g_k(k)$$$.
Consider difference method, $$$h_{k,n}-h_{k,n-1}=x^{n+1}g_k(n)-x^ng_k(n-1)=n^kx^n$$$.
So $$$xg_k(n)-g_k(n-1)=n^k$$$. But we still can't get $$$C$$$. There is a conclusion: the $$$(n + 1)$$$-th difference of the $$$n$$$-degree polynomial $$$f(x)$$$ is $$$0$$$. That is to say, $$$\sum_{i=0}^{n+1}(-1)^{n+1-i}\binom{n+1}{i}f(i)=0$$$.
So we can get $$$g_k(0)$$$, then use polynomial fast interpolation, and then polynomial multipoint evaluation to get the value of every point we need.
It can be observed that the DP transfer can be replaced by matrix, so there are matrices each position, the problem is equivalent to get the sum of the matrix products of all intervals. It's easy to solve. You can solve it in the way you like.
Overall complexity: $$$\mathcal{O}(k\log^2k+n)$$$.