题解 P5591 【小猪佩奇学数学】

· · 题解

题意

给定 n,\,p,\,k ，求：

\sum\limits_{i=0}^{n}\binom{n}{i}\,p^i\left\lfloor\frac{i}{k}\right\rfloor

解题思路

看到题目给出的式子第一反应是什么？二项式定理！

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

移掉烦人的 \left\lfloor \frac{i}{k} \right\rfloor ，二项式定理告诉我们，上面的式子等于 (p + 1)^n 。

然而我们无法轻易地将 \left\lfloor \frac{i}{k} \right\rfloor 表示成某个数的 i 次幂，不妨先来解决子问题。令 k = 1 ：

\sum\limits_{i=0}^{n}\binom{n}{i}\,p^i\left\lfloor\frac{i}{k}\right\rfloor = \sum\limits_{i=0}^{n}\binom{n}{i}\,p^i\,i

因为 n 很大，如果不将和式转换成封闭形式，绝对会超时。而二项式定理，就是一种关键的途径。

所以我们希望能把孤零零的 i 去掉，或者放在某个 i 次幂里。\binom{n}{i}i 给予了很大的机会。

这里有一个引理：

\binom{n}{m}m = \binom{n-1}{m-1}n

刚刚好 \binom{n}{i}i = \binom{n-1}{i-1}n ，消除了 i 。接下来就好办很多了：

\sum\limits_{i=0}^{n}\binom{n}{i}\,p^i\,i = np\sum\limits_{i=0}^{n}\binom{n-1}{i-1}\,p^{i-1}\ = np(p+1)^{n-1}

然后我们再把目光移到一般情况上。但还是老问题，\left\lfloor \frac{i}{k} \right\rfloor 如何表示成某个数的 i 次幂呢？它无法像上面一样与 \binom{n}{i} 有关系。注意到，k 是很小的 2 次幂，模数 9982444353 = 199\times2^{23}+1 ，这提示着，我们需要单位根反演。不难发现：

\left\lfloor \frac{i}{k} \right\rfloor = (\sum_{j=0}^{i}\,[k\,|\,j]) - 1

而单位根反演的结论恰好是：

[k\,|\,n] = \frac{1}{k}\sum_{d=0}^{k-1}\omega_{k}^{dn}

代入得：

\left\lfloor \frac{i}{k} \right\rfloor = (\sum_{j=0}^{i}\,[k\,|\,j]) - 1 = (\sum_{j=0}^{i}\frac{1}{k}\sum_{d=0}^{k-1}\omega_{k}^{dj}) - 1

挪到原式中并化简：

\begin{aligned}\sum\limits_{i=0}^{n}\binom{n}{i}\,p^i\left\lfloor\frac{i}{k}\right\rfloor &= (\sum\limits_{i=0}^{n}\binom{n}{i}\,p^i\sum_{j=0}^{i}\frac{1}{k}\sum_{d=0}^{k-1}\omega_{k}^{dj}) - (\sum\limits_{i=0}^{n}\binom{n}{i}\,p^i) \\ &= (\sum\limits_{i=0}^{n}\binom{n}{i}\,p^i\sum_{j=0}^{i}\frac{1}{k}\sum_{d=0}^{k-1}\omega_{k}^{dj}) - (p+1)^n && (\text{后面有二项式定理}) \\ &= (\frac{1}{k}\sum_{d=0}^{k-1}\sum\limits_{i=0}^{n}\binom{n}{i}\,p^i\sum_{j=0}^{i}\omega_{k}^{dj}) - (p+1)^n && (\text{这一步很关键}) \\ &= (\frac{1}{k}\sum_{d=0}^{k-1}\sum\limits_{i=0}^{n}\binom{n}{i}\,p^i\frac{\omega_{k}^{d(i+1)}-1}{\omega_{k}^{d}-1}) - (p+1)^n && (\text{等比数列求和}) \\ &= (\frac{1}{k}\sum_{d=0}^{k-1}\frac{\sum\limits_{i=0}^{n}\binom{n}{i}\,p^i(\omega_{k}^{d(i+1)}-1)}{\omega_{k}^{d}-1}) - (p+1)^n && (\text{与 i 无关的分母，提取出}) \\ &= (\frac{1}{k}\sum_{d=0}^{k-1}\frac{\omega_{k}^{d}(\sum\limits_{i=0}^{n}\binom{n}{i}\,p^i(\omega_{k}^d)^i)-(\sum\limits_{i=0}^{n}\binom{n}{i}\,p^i)}{\omega_{k}^{d}-1}) - (p+1)^n && (\text{拆开}) \\ &= (\frac{1}{k}\sum_{d=0}^{k-1}\frac{\omega_{k}^{d}(p\omega_k^{d}+1)^n-(p+1)^{n}}{\omega_{k}^{d}-1}) - (p+1)^n && (\text{舒服！两边都可以二项式定理})\end{aligned}

于是 O(k \log) 的时间内就能算出答案啦。根据单位根定义，\omega_{k}^{d} = (\omega_{k}^{1})^d = (g^{(mod-1)/k})^d ，其实 g 是模数 998244353 的原根，取 3 即可。

然后实现代码，发现样例都过不去。名场面？

问题出在哪儿呢？重新看看等比数列求和？当 d = 0,\,\omega_{k}^{d} = 1 ，分母不就为 0 了？

我们需要特判 d = 0 。

\sum\limits_{i=0}^{n}\binom{n}{i}\,p^i\sum_{j=0}^{i}\omega_{k}^{dj} = \sum\limits_{i=0}^{n}\binom{n}{i}\,p^i(i+1) = (\sum\limits_{i=0}^{n}\binom{n}{i}\,p^ii) + (\sum\limits_{i=0}^{n}\binom{n}{i}\,p^i)

运气真好，推着推着又回到了热身问题。根据开始的结论，它等于 np(p+1)^{n-1} + (p+1)^n 。

至此，本题已被完美地解决。时间复杂度 O(k \log) 。

代码实现

#include <bits/stdc++.h>

const int mod = 998244353, gen = 3;

inline int add(int x, int y) { return x + y >= mod ? x + y - mod : x + y; }
inline int sub(int x, int y) { return x - y >= 0 ? x - y : x - y + mod; }
inline int power(int x, int y, int res = 1) {
    for (; y; y >>= 1, x = 1ll * x * x % mod) {
        if (y & 1) { res = 1ll * res * x % mod; }
    } return res;
}

const int N = 2e6 + 5;

int n, p, k, ans, w[N];

int main() {
    scanf("%d%d%d", &n, &p, &k);
    w[0] = 1; w[1] = power(gen, (mod - 1) / k);
    for (int i = 2; i < k; i++) { w[i] = 1ll * w[i - 1] * w[1] % mod; }
    ans = add(1ll * n * p % mod * power(add(p, 1), n - 1) % mod, power(add(p, 1), n));
    for (int i = 1; i < k; i++) {
        int now = 1ll * w[i] * power(add(1ll * w[i] * p % mod, 1), n) % mod;
        now = sub(now, power(add(p, 1), n));
        ans = add(ans, 1ll * now * power(sub(w[i], 1), mod - 2) % mod);
    }
    ans = sub(1ll * ans * power(k, mod - 2) % mod, power(add(p, 1), n));
    printf("%d\n", ans);
    return 0;
}