洛谷P5907 数列求和加强版 / SPOJ MOON4

题面描述

求

$$\sum_{i=1}^ni^ka^i$$

对 $998244353$ 取模的结果。

$O(k)$做法

为了推导方便，下令 $p = a^{-1}$。即求

$$\sum_{i=1}^n\frac{i^k}{p^i}$$

我们裂项，即尝试寻找多项式 $f(x)$，使得

$$\frac{x^k}{p^x} = \frac{f(x)}{p^{x-1}} - \frac{f(x+1)}{p^x}$$

即

$$x^k = pf(x) - f(x+1)(*)$$

答案将是 $f(1) - \frac{f(n+1)}{p^n}$。

显然 $\deg f = k$。

这样临项点值的差，由于恒等式 $(x + 1)^{\underline{n}}- x ^ {\underline{n}} = nx^{\underline{n-1}}$，可以考虑展开成下降幂来推导。具体地，设 $f(x) = \sum\limits_{i = 0}^n c_ix^{\underline{i}}$，则可以展开 $(*)$ 式：

$$\begin{aligned} x^k = \sum^k_{i = 0}\begin{Bmatrix}k \\ i\end{Bmatrix}x^{\underline{i}} &= \sum_{i = 0}^k pc_ix^{\underline{i}} - \sum_{i = 0}^k c_i(x + 1)^{\underline{i}} \\ &= \sum_{i = 0}^k(p-1)c_ix^{\underline{i}} - \sum_{i=1}^kic_ix^{\underline{i-1}} \\ &= (p-1)c_kx^{\underline{k}}+\sum_{i=0}^{k-1}((p-1)c_i-(i+1)c_{i+1})x^{\underline{i}} \end{aligned}$$

对比系数能够得到：

$$(p - 1)c_k = \begin{Bmatrix}k \\ k\end{Bmatrix} = 1$$

$i = 0, 1...,k-1$：

$$(p - 1)c_i - (i+1)c_{i+1} = \begin{Bmatrix}k \\ i\end{Bmatrix}$$

求一行斯特林数是 $O(k \log k)$的，需要进一步优化。

直接求每一项系数也许有些困难，但是这里有一个取巧的办法：
如果我们知道 $f(0)$，那么通过$(*)$式，可以递推出 $f(1), ...f(k)$，那么就可以利用这 $k+1$ 个点值进行线性拉格朗日插值，避开了这个难处。

而 $f(0)=c_0$，因此下面我们只需要求出 $c_0$ 即可。省略这一步过程，我们恰好可以将这个系数的递推式求出来，即

$$c_0 = \frac 1 {p - 1}\sum_{i = 0}^k\begin{Bmatrix}k \\ i\end{Bmatrix} \frac{i!}{(p-1)^i}$$

这个式子是可以 $O(k)$ 求解的，参见 EI 的 Binomial Sum。由于我也是第一次学这个所以下面还是写一下过程。

下面令 $t = (p-1)^{-1}$ ，求

$$\sum_{i = 0}^k\begin{Bmatrix}k \\ i\end{Bmatrix} t^i i!$$

写成生成函数就是

$$\begin{aligned} \sum_{i = 0}^kt^i \begin{Bmatrix}k \\ i\end{Bmatrix} i! &= \sum_{i = 0}^k t^i \left[\frac{z^k}{k!}\right](\exp z-1)^i \\ &= \left[\frac{z^k}{k!}\right]\sum_{i=0}^k (t(\exp z - 1))^i \\ &= \left[\frac{z^k}{k!}\right] \frac 1 {1 - t(\exp z - 1)} \end{aligned}$$

令 $u = \exp z$，对于 $z$ 来说，这本应是一个无穷级数，但利用 $\exp z - 1$ 没有常数项的性质，我们可以将对于 $u$ 的求和上限限制在 $k$ 处（截断），之后的项都不会有贡献。因此就是求

$$\left[\frac{z^k}{k!}\right] \frac {1 - (t(\exp z - 1))^{k+1}} {1 - t(\exp z - 1)}$$

那么 GF 写作：

$$G(u) = \frac {1 - (t(u-1))^{k+1}}{1-t(u-1)}$$

这是一个关于 $u$ 的有理分式，将分母移到一边就可以 $O(k)$ 求出 $u^i$ 处系数，然后就做完了。退役OIer不想省略下面的细节，继续推导。我们设 $G(u) = \sum\limits_{i\ge0}g_iu^i$，那么

$$g_0 = G(0) = \frac{1+t^{k+1}(-1)^k}{1+t}$$

$$(1+t-tu)G = 1-t^{k+1}(u-1)^{k+1}$$

提取系数就是

$$(1+t)g_i-tg_{i-1}=t^{k+1}\binom{k+1}{i}(-1)^{k-i}$$

最后，我们得到了，

$$\begin{aligned} c_0 &= t \left[\frac{z^k}{k!}\right] \sum_{i=0}^k g_i \exp iz \\ &= t\sum_{i=0}^k g_i i^k \end{aligned}$$

这道题几乎做完了。剩下的就是特判 $a=p=1$ 的情况，此时问题退化为自然数幂和，可以直接线性拉格朗日插值。

代码（代码中，使用 $m$ 代替了 $k$）：

#include <cstdio>
#define ll long long
const int M = 1e7, mod = 998244353;
inline int mul(int a, int b) {return (ll)a * b % mod;}
inline int add(int a, int b) {int ret = a + b; return ret >= mod ? ret - mod : ret;}
inline int minus(int a, int b) {int ret = a - b; return ret < 0 ? ret + mod : ret;}
inline int qpow(int a, int b) {
    int ret = 1;
    for(; b; b >>= 1, a = mul(a, a))
        if(b & 1)   ret = mul(ret, a);
    return ret;
}
int f[M + 5], g[M + 5], p, n, m, ip, ip1;
int fac[M + 5], invf[M + 5];
int pre, suf[M + 5];
int cnt, v[M + 5], pr[M >> 2], pwm[M + 5];
inline int C(int m, int n) {return (ll)fac[m] * invf[n] % mod * invf[m - n] % mod;}
inline int pn1(int num, int p) {return p & 1 ? mod - num : num;}

int main() {
    scanf("%d%d%d", &n, &ip, &m);
    //initialization
    fac[0] = 1;
    p = qpow(ip, mod - 2);
    for(int i = 1; i <= m + 1; ++i) fac[i] = mul(fac[i - 1], i);
    invf[m + 1] = qpow(fac[m + 1], mod - 2);
    for(int i = m; i >= 0; --i) invf[i] = mul(invf[i + 1], i + 1);

    pwm[1] = 1;
    for(int i = 2; i <= m + 2; ++i) {
        if(!v[i]) v[i] = pr[++cnt] = i, pwm[i] = qpow(i, m);
        for(int j = 1; j <= cnt; ++j) {
            if(pr[j] > v[i] || i * pr[j] > m + 2)    break;
            v[i * pr[j]] = pr[j];
            pwm[i * pr[j]] = mul(pwm[i], pwm[pr[j]]);
        }
    }
    ip1 = qpow(p - 1, mod - 2);
    if(p ^ 1) {
	    //calulate g, f
	    int t1 = qpow(ip1 + 1, mod - 2), t2 = qpow(ip1, m + 1);
	    // t1 = 1/t + 1, t2 = t ^ {m + 1}
	    g[0] = mul(t1, 1 + pn1(t2, m));
	    for(int i = 1; i <= m; ++i) {
	        g[i] = mul(t1, mul(ip1, g[i - 1]) + mul(C(m + 1, i), pn1(t2, m - i)));
	        f[0] = add(f[0], mul(pwm[i], g[i]));
	    }
		f[0] = mul(ip1, f[0]);
	}
    if(p ^ 1)	for(int i = 1; i <= m; ++i)
        	f[i] = minus(mul(p, f[i - 1]), pwm[i - 1]);
	else for(int i = 1; i <= m + 1; ++i)
		f[i] = add(f[i - 1], pwm[i]);
       
    //prepare for lagruange interpolation
    if(p == 1)	--n, ++m;
    pre = n + 1, suf[m + 1] = 1;
//    for(int i = 1; i <= m; ++i) pre[i] = mul(pre[i - 1], (mod + n + 1 - i) % mod);
    for(int i = m; i >= 1; --i) suf[i] = mul(suf[i + 1], minus(n + 1, i));
    
    //calculate answer
    int Ans = mul(f[0], pn1(mul(suf[1], invf[m]), m));
    for(int i = 1; i <= m; ++i)
       Ans = add(Ans, pn1(mul(f[i], mul(mul(pre, suf[i + 1]), mul(invf[i], invf[m - i]))), m - i)),
    	pre = mul(pre, minus(n + 1, i));
    if(p ^ 1)	Ans = minus(f[1], mul(Ans, qpow(qpow(p, n), mod - 2)));
	
    printf("%d", Ans);
    return 0;
}