题意概括
给出两个序列 \(a_0,a_2,\cdots a_{n-1}\) ,\(b_0,b_2,\cdots b_{m-1}\) ,从两个序列中各等概率的选出两个数 \(a_i,b_j\) ,对于 \(k\in [1,t]\) ,分别求出 \((a_i+b_j)^k\) 在 \(mod\ 998244353\) 意义下的期望值。
\(n,m\le 1e5;\ \ t\le 1e5;\ \ a_i,b_i< 998244353\)
符号规定
下文使用形如 \(f(x)\) 的式子表示多项式, \(f_p(x)\) 表示 \(f(x)\) 的 \(x^p\) 项的系数。\(f*g\) 表示 \(f(x),g(x)\) 的卷积,\((f*g)_k\) 表示卷积后 \(x^k\) 项系数。
思路分析
要求期望权值,可以考虑算出所有情况的权值和再除以情况数,因此接下来我们只考虑如何快速求所有情况的权值和。
对于每一个 \(k\) ,权值和为 \(\sum_{i=0}^{n-1}\sum_{j=0}^{m-1}(a_i+b+j)^k\) ,考虑进行二项式展开,得到:
\[\begin{aligned}
& \sum_{i=0}^{n-1}\sum_{j=0}^{m-1}\sum_{p=0}^{k}C_{k}^{p}\times a_i^{p}\times b_j^{k-p}\\
& \sum_{i=0}^{n-1}\sum_{j=0}^{m-1}\sum_{p=0}^{k}\frac{k!}{p!\times(k-p)!}\times a_i^p\times b_j^{k-p}\\
& k!\times \sum_{p=0}^{k}\sum_{i=0}^{n-1}\frac{a_i^p}{p!}\sum_{j=0}^{m-1}\frac{b_j^{k-p}}{(k-p)!}\\
& k!\times \sum_{p=0}^{k}(\sum_{i=0}^{n-1}\frac{a_i^p}{p!})(\sum_{j=0}^{m-1}\frac{b_j^{k-p}}{(k-p)!})
\end{aligned}
\]
但是可以发现如果每一次都去枚举 \(k\) 的值的话,复杂度至少为 \(O(nk)\) ,并且当前的形式类似卷积,所以考虑 \(NTT\) ,定义两个多项式 \(f,g\) 并满足:
\[f_p(x)=\sum_{i=0}^{n-1}\frac{a_i^p}{p!},g_p(x)=\sum_{i=0}^{m-1}\frac{b_i^p}{p!}
\]
记 \(t=k\) 时的权值和为 \(ans(x)\) 的 \(x^k\) 项的系数,那么就有优美的卷积:
\[ans_k(x)=k!\times(f*g)_k
\]
这一步卷积时间复杂度 \(O(n\log n)\) ,所以现在的问题在于快速求 \(f(x),g(x)\) 。因为 \(f(x),g(x)\) 是类似的,所以下面就只探讨如何求解 \(f(x)\) 。将 \(f(x)\) 展开来写,得到:
\[\begin{aligned}
f(x)&=\sum_{i=0}x^i\sum_{j=0}^{n-1}\frac{a_j^i}{i!}\\
&=\sum_{i=0}\frac{x^i}{i!}\sum_{j=0}^{n-1}a_j^i
\end{aligned}
\]
可以发现 \(f(x)\) 的第 \(i\) 项的系数 \(\frac{1}{i!}\) 可以先不考虑,最后再统一处理,因此先不管 \(\frac{1}{i!}\) ,考虑优化计算 \(h(x)=\sum_{i=0}x^i\sum_{j=0}^{n-1}a_j^i\) ,可以用生成函数优化:
\[\begin{aligned}
h(x)&=\sum_{i=0}x^i\sum_{j=0}^{n-1}a_j^i\\
&=\sum_{i=0}^{n-1}\sum_{j=0}a_i^jx^j\\
&=\sum_{i=0}^{n-1}\frac{1}{1-a_ix}
\end{aligned}
\]
加法计算复杂度过高,可以考虑用 \(\exp\) 加法转乘法,因为发现 \(h(x)=\ln'(1-a_ix)\) ,所以可以考虑先构造一个多项式 \(F(x)=\prod_{i=0}^{n-1}(1-a_ix)\) ,这个多项式是可以 分治NTT 求的,复杂度 \(O(n\log^2n)\) ,然后令 \(G(x)=\ln F(x)\) ,则:
\[G(x)=\sum_{i=0}^{n-1}\ln(1-a_ix)
\]
这样就出现了我们想要的形式,但其实这样做完之后还没完全结束,还要转换一步:
\[\begin{aligned}
G'(x)&=\sum_{i=0}^{n-1}(\ln(1-a_ix))'\\
&=\sum_{i=0}^{n-1}(\ln'(1-a_ix)\times (1-a_ix)')\\
&=\sum_{i=0}^{n-1}\frac{-a_i}{1-a_ix}\\
&=-\frac{1}{x}\sum_{i=0}^{n-1}\frac{a_ix}{1-a_ix}\\
&=-\frac{1}{x}\sum_{i=0}^{n-1}\frac{1-(1-a_ix)}{1-a_ix}\\
&=-\frac{1}{x}\sum_{i=0}^{n-1}(\frac{1}{1-a_ix}-1)\\
&=-\frac{1}{x}(h(x)-n)\\
\end{aligned}
\]
因此, \(h(x)=-xG(x)+n\) 。到这里这题就结束了,总结一下就是:
\[\begin{aligned}
&F(x)=\prod_{i=0}^{n-1}(1-a_ix)\\
&G(x)=\ln(F(x))\\
&h(x)=-xG'(x)+n\\
&f(x)=\sum_{i=0}^{t}\frac{h_i(x)}{i!}\times x^i\\
&ans(x)=\sum_{i=0}^{t}(f*g)_i\times i!\times x^i
\end{aligned}
\]
最后当 \(t=k\) 时,输出 \(\frac{ans_t(x)}{nm}\) 即可。下面是 代码 。