P5591 小猪佩奇学数学 题解

· · 题解

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

有：

\left\lfloor \frac{i}{k} \right\rfloor = \frac{i - i \bmod k}{k}

所以有：

\begin{aligned} & \sum \limits_{i=0}^n \binom{n}{i} \times p^i \times \left\lfloor \frac{i}{k} \right\rfloor\\ = & \frac{1}{k} \sum \limits_{i=0}^n \binom{n}{i} \times p^i \times (i - i \bmod k)\\ \end{aligned}

先省略前面的 \frac{1}{k}：

\begin{aligned} & \sum \limits_{i=0}^n \binom{n}{i} \times p^i \times (i - i \bmod k)\\ = & \sum \limits_{i=0}^n \binom{n}{i} p^i i - \sum \limits_{i=0}^n \binom{n}{i} p^i (i \bmod k)\\ \end{aligned}

式子左边有：

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

式子右边有：

\begin{aligned} & \sum \limits_{i=0}^n \binom{n}{i} p^i (i \bmod k)\\ = & \sum \limits_{i=0}^n \binom{n}{i} p^i \sum \limits_{j=0}^{k-1} j [i \bmod k = j]\\ = & \sum \limits_{i=0}^n \binom{n}{i} p^i \sum \limits_{j=0}^{k-1} j \frac{1}{k} \sum \limits_{s=0}^{k-1} \omega_k^{is} \omega_k^{-js}\\ = & \frac{1}{k} \sum \limits_{s = 0}^{k-1} \sum \limits_{i=0}^n \binom{n}{i} p^i \omega_k^{is} \sum \limits_{j=0}^{k-1} j \omega_k^{-js}\\ = & \frac{1}{k} \sum \limits_{s=0}^{k-1} \left( \sum \limits_{i=0}^n \binom{n}{i} (p \omega_k^s)^i \right) \sum \limits_{j=0}^{k-1} j \omega_k^{-js}\\ = & \frac{1}{k} \sum \limits_{s=0}^{k-1} (p \omega_k^s + 1)^n \sum \limits_{j=0}^{k-1} j \omega_k^{-js}\\ \end{aligned}

现在设法求出 \sum \limits_{j=0}^{k-1} j \omega_k^{-js}。

设 f(k, p) = \sum \limits_{j=0}^k j p^j。

已知：

\sum \limits_{j=0}^k p^j = \frac{p^{k+1} - 1}{p - 1}

所以有：

\begin{aligned} p f(k, p) + \sum \limits_{j=0}^k p^j - 1 - k p^{k+1} & = f(k, p)\\ p f(k, p) + \frac{p^{k+1} - 1}{p - 1} - 1 - k p^{k+1} & = f(k, p)\\ \frac{p^{k+1} - 1}{p - 1} - 1 - k p^{k+1} & = f(k, p) (1 - p)\\ f(k, p) & = \frac{\frac{p^{k+1} - 1}{p - 1} - 1 - k p^{k+1}}{1 - p} \end{aligned}

注意，当 p = 1 时需要特判 f(k, p) = \frac{k (k+1)}{2}。

所以：

\begin{aligned} & \sum \limits_{j=0}^{k-1} j \omega_k^{-js}\\ = & f(k-1, \omega_k^{-s})\\ \end{aligned}

所以右式：

\begin{aligned} & \frac{1}{k} \sum \limits_{s=0}^{k-1} (p \omega_k^s + 1)^n \sum \limits_{j=0}^{k-1} j \omega_k^{-js}\\ = & \frac{1}{k} \sum \limits_{s=0}^{k-1} (p \omega_k^s + 1)^n f(k-1, \omega_k^{-s})\\ \end{aligned}

所以原式：

\begin{aligned} & \frac{1}{k} \left( \sum \limits_{i=0}^n \binom{n}{i} p^i i - \sum \limits_{i=0}^n \binom{n}{i} p^i (i \bmod k) \right)\\ = & \frac{1}{k} \left( np (p + 1) ^{n-1} - \frac{1}{k} \sum \limits_{s=0}^{k-1} (p \omega_k^s + 1)^n f(k-1, \omega_k^{-s})\right) \end{aligned}

代码并不难写。

#include <iostream>
#include <cstdio>
using namespace std;
typedef long long ll;

const ll mod = 998244353;

inline ll chk(ll x) { return x >= mod ? x-mod : x; }

inline ll fpm(ll a, ll k = mod-2) {
    ll res = 1;
    while(k) {
        if(k&1) res = res*a % mod;
        a = a*a % mod;
        k >>= 1;
    }
    return res;
}

inline ll f(ll k, ll p) {
    if(p == 1) return k * (k+1) / 2 % mod;
    ll res = chk(mod + fpm(p, k+1) - 1) * fpm(p-1) % mod;
    res = chk(res + mod - 1);
    res = chk(res + 1ll * (mod-1) * k % mod * fpm(p, k+1) % mod);
    return 1ll * res * fpm(mod+1-p) % mod;
}

int n, p, k, w, ans;

int main() {
    scanf("%d %d %d", &n, &p, &k);
    w = fpm(3, (mod-1) / k);

    for(int s = 0; s < k; s++)
        ans = chk(ans + 1ll * fpm(1ll * p * fpm(w, s) % mod + 1, n) * f(k-1, fpm(fpm(w), s)) % mod);
    ll res = ans = 1ll * ans * fpm(k) % mod * (mod-1) % mod;
    ans = chk(ans + 1ll * n * p % mod * fpm(p+1, n-1) % mod);
    ans = 1ll * ans * fpm(k) % mod;

    printf("%d\n", ans);
    return 0;
}