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

单位根反演套路，记一下。

思路

单位根反演。

除法下取整不能整除分块，只能转化一下。考虑将原式转化成 \(\sum\limits_{i = 0}^n {n \choose i} p^i \frac{i - i \bmod k}{k}\).

考虑拆成两个和式：\(\sum\limits_{i = 0}^n {n \choose i} p^i i - \sum\limits_{i = 0}^n {n \choose i} p^i (i \bmod k)\).

结论：\({n \choose m} m = {n - 1 \choose m - 1} n\).

于是第一部分为 \(np \sum\limits_{i = 0}^n {n - 1 \choose i - 1} p^{i - 1}\)，二项式定理得 \(np (p + 1)^{n - 1}\).

第二部分可以考虑枚举余数：\(\sum\limits_{i = 0}^n {n \choose i} p^i i - \sum\limits_{i = 0}^n {n \choose i} p^i \sum\limits_{j = 0}^{k - 1} j [i \equiv j \pmod k]\).

单位根反演得 \(\sum\limits_{i = 0}^n {n \choose i} p^i i - \sum\limits_{i = 0}^n {n \choose i} p^i \sum\limits_{j = 0}^{k - 1} j \frac{1}{k} \sum\limits_{x = 0}^{k - 1} w_{k}^{(i - j) x}\).

简单地交换求和顺序并化简得：\(\sum\limits_{j = 0}^{k - 1} j \sum\limits_{t = 0}^{k - 1} w_{k}^{-tj} (1 + p w_k^t)^n\).

现在得到了 \(O(k^2)\) 做法，但是注意到单位根有性质，可以进一步优化复杂度。

令 \(S(n, k) = \sum\limits_{i = 0}^{n - 1} i k^i\).

简单求和得 \(S(n, k) = \frac{\frac{1 - k^n}{k - 1} + (n - 1) k^n + 1}{k - 1}\).

注意到 \(n\) 次单位根的性质 \(w_k^n = 1\)，代入的 \(k\) 总是 \(n\) 次单位根。

所以原式化简成 \(S(n, k) = \frac{n}{k - 1}\).

注意 \(k = 1\) 时特判 \(S(n, 1) = \frac{n (n - 1)}{2}\).

最终原式的形式是 \(\sum\limits_{t = 0}^{k - 1} (1 + p w_k^t)^n S(k - 1, w_k^{-t})\).

代码

#include <cstdio>
using namespace std;

typedef long long ll;

const int maxk = 1.1e6 + 5;
const int mod = 998244353;
const int g = 3;

int n, p, k;
ll w[maxk];

ll qpow(ll base, ll power = mod - 2)
{
    ll res = 1;
    while (power)
    {
        if (power & 1) res = res * base % mod;
        base = base * base % mod, power >>= 1;
    }
    return res;
}

ll S(ll n, ll k)
{
    if (k == 1) return ((n - 1) * n >> 1) % mod;
    return n * qpow(k - 1) % mod;
}

int main()
{
    scanf("%d%d%d", &n, &p, &k);
    ll r = qpow(3, (mod - 1) / k), ans = 0;
    w[0] = 1;
    for (int i = 1; i <= k; i++) w[i] = w[i - 1] * r % mod;
    for (int i = 0; i < k; i++) ans = (ans + qpow((p * w[i] + 1) % mod, n) * S(k, w[k - i])) % mod;
    ans = ans * qpow(k) % mod;
    ans = (1ll * n * p % mod * qpow(p + 1, n - 1) - ans + mod) % mod;
    ans = ans * qpow(k) % mod;
    printf("%lld\n", ans);
    return 0;
}