【小猪佩奇学数学 题解】一定要写一些东西才看起来可爱嘛

· · 题解

计算：

\sum_{i=0}^n \binom n i p^{i} \left\lfloor \dfrac{i}{k} \right\rfloor \bmod 998244353

其中 1 \leq n,p <998244353,k \in \{2^{w}|0 \leq w \leq 20\}。

把 \left\lfloor \dfrac{i}{k} \right\rfloor 拆为 \dfrac{i-(i\bmod k)}{k}。

先看

\sum_{i=0}^n \binom n i p^{i} i

由于

\binom n i=\dfrac{n!}{i!(n-i)!}=\dfrac n i\binom {n-1}{i-1}

所以有

\begin{aligned} &\sum_{i=0}^n \binom n i p^{i} i\\ =&n\sum_{i=0}^{n-1} \binom {n-1} {i} p^{i+1}\\ =&np\sum_{i=0}^{n-1} \binom {n-1} {i} p^i\\ =&np(p+1)^n \end{aligned}

再看另一部分

\sum_{i=0}^n \binom n i p^{i} (i\bmod k)=\sum_{i=0}^{k-1}i\sum_{j\equiv i\pmod k}\binom n jp^j

可以发现

[x^j](px+1)^n=\binom n jp^j

所以只需要计算这个多项式时，把次数取模即可。

FFT 有一个性质是：当多项式乘出来的长度大于进行 FFT 的长度 n 时，次数大于等于 n 的项 i，系数会加到 i\bmod n 上。也就是说，FFT 是循环的。

由于这一题 k 为 2 的次幂，所以可以直接运用上面这条结论，先对多项式 px+1 进行长度为 k 的 NTT，再对每个位置 n 次方，最后 INTT，此时第 i 项就是 \sum\limits_{j\equiv i\pmod k}\binom n jp^j 的值。

将两部分的结果相减，除以 k 就是最终答案。

时间复杂度 O(k\log n)。

#include <bits/stdc++.h>

using namespace std;

typedef long long LL;

const int N = 1 << 20 | 5, P = 998244353, G = 3;

int n, p, k;
LL a[N];

inline LL ksm(LL a, LL n)
{
    LL res = 1;
    while(n)
    {
        if(n & 1) res = res * a % P;
        a = a * a % P;
        n >>= 1;
    }
    return res;
}
inline void NTT(LL a[], int n, int ty)
{
    for(int i = 0, j = 0; i < n; i ++ )
    {
        if(i < j) a[i] ^= a[j] ^= a[i] ^= a[j];
        for(int k = (n >> 1); (j ^= k) < k; k >>= 1) ;
    }
    static LL w[N];
    w[0] = 1;
    for(int i = 1; i < n; i <<= 1)
    {
        LL wi = ksm(G, P - 1 + ty * (P - 1) / (i << 1));
        for(int j = i - 2; j >= 0; j -= 2)
        {
            w[j] = w[j >> 1];
            w[j | 1] = w[j] * wi % P;
        }
        for(int j = 0; j < n; j += (i << 1))
        {
            for(int k = j; k < j + i; k ++ )
            {
                LL t0 = a[k], t1 = a[k + i] * w[k - j] % P;
                (a[k] = t0 + t1) >= P && (a[k] -= P), (a[k + i] = t0 - t1 + P) >= P && (a[k + i] -= P);
            }
        }
    }
    if(ty == -1)
    {
        LL inv = ksm(n, P - 2);
        for(int i = 0; i < n; i ++ ) a[i] = a[i] * inv % P;
    }
}

int main()
{
    scanf("%d%d%d", &n, &p, &k);
    LL res = ksm(p + 1, n - 1) * n % P * p % P;
    a[0] = 1, a[1] = p;
    NTT(a, k, 1);
    for(int i = 0; i < k; i ++ ) a[i] = ksm(a[i], n);
    NTT(a, k, -1);
    for(int i = 1; i < k; i ++ ) res = (res + (P - i) * a[i]) % P;
    printf("%lld\n", res * ksm(k, P - 2) % P);
    return 0;
}