P5591 小猪佩奇学数学 题解

（好久没写过某一道题的题解了）

一道牛牛的推式子题。

---

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

\[\sum_{i=0}^n{n\choose i}\times p^i\times \lfloor\frac i k\rfloor \]

对 \(998244353\) 取模。\(n,p < 998244353,k\in\{2^w|0\le w\le 20\}\)

前置知识：单位根反演。定义 \(\omega_n\) 表示复平面单位圆上的 \(n\) 次单位根，有 \(\omega_n^n=1\)。单位根反演为：

\[[n|m] = \frac 1 n\sum_{i=0}^{n-1}\omega_n^{im} \]

证明：

• 若 \(n|m\)，那么 \(\omega_n^m=1\)，于是和式每一项都为 \(1\)，\(\frac 1 n\sum\limits_{i=0}^{n-1}\omega_n^{im} = \frac 1 n\times n = 1\)。
• 若 \(n\nmid m\)，那么根据等比数列求和公式：\(\sum\limits_{i=0}^{n-1}\omega_n^{im} = \frac{\omega_n^{nm}-1}{\omega_n^m-1} = 0\)

于是当我们遇到和式中有 \([n|m]\) 的式子时，就可以用单位根反演，然后交换和式进行操作。

一般来说，如果要用单位根反演且有模数，那么要使 \(\omega_n\) 在模意义下存在，设 \(a\) 为模数 \(p\) 的一个原根，那么 \(\omega_n\equiv a^\frac{p-1}n\pmod p\)。

所以一般题面都会保证 \(n\mid p-1\)，例如模数为 \(998244353\) 时，要求 \(n\) 为不超过 \(20\) 的幂次（本题）。反过来，当你看到类似的提示时，你应该反应过来这是单位根反演。

---

本题：

因为下取整非常不好表示，所以可以转化为 \(\lfloor\frac n m\rfloor = (\sum\limits_{i=0}^n[m|i])-1\)，然后考虑用单位根反演。

\[\begin{aligned} \sum_{i=0}^n{n\choose i}\times p^i\times \lfloor\frac i k\rfloor &= \sum_{i=0}^n{n\choose i}\times p^i\times(\sum_{j=0}^i[k|j]-1) \\ &= \sum_{i=0}^n{n\choose i}\times p^i\times\sum_{j=0}^i\frac 1 k\sum_{d=0}^{k-1}\omega_k^{jd}-\sum_{i=0}^n{n\choose i}\times p^i \\ &= \frac 1 k\sum_{d=0}^{k-1}\sum_{i=0}^n{n\choose i}\times p^i\times\sum_{j=0}^i\omega_k^{jd}-(p+1)^n \\ &= \frac 1 k\sum_{d=1}^{k-1}\sum_{i=0}^n{n\choose i}\times p^i\times\sum_{j=0}^i\omega_k^{jd}+\frac 1 k\sum_{i=0}^n{n\choose i}\times p^i\times (i+1)-(p+1)^n \\ &= \frac 1 k\sum_{d=1}^{k-1}\sum_{i=0}^n{n\choose i}\times p^i\times\frac{\omega_k^{(i+1)d}-1}{\omega_k^d-1}+\frac 1 k(\sum_{i=0}^n{n\choose i}\times p^i\times i+(p+1)^n)-(p+1)^n \\ &= \frac 1 k\sum_{d=1}^{k-1}\frac{\sum\limits_{i=0}^n{n\choose i}\times p^i\times(\omega_k^{di}\times \omega_k^d-1)}{\omega_k^d-1}+\frac 1 k(\sum_{i=0}^n n\times {n-1\choose i-1}\times p^{i-1}\times p+(p+1)^n)-(p+1)^n \\ &= \frac 1 k\sum_{d=1}^{k-1}\frac{\omega_k^d\times\sum\limits_{i=0}^n{n\choose i}\times p^i\times\omega_k^{di}-\sum\limits_{i=0}^n{n\choose i}\times p^i}{\omega_k^d-1}+\frac 1 k(np(p+1)^{n-1}+(p+1)^n)-(p+1)^n \\ &= \frac 1 k\sum_{d=1}^{k-1}\frac{\omega_k^d\times(p\times \omega_k^d+1)^n-(p+1)^n}{\omega_k^d-1}+\frac 1 k(np(p+1)^{n-1}+(p+1)^n)-(p+1)^n \end{aligned} \]

至此，可以 \(\mathcal{O}(k\log n)\) 做。

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#define ll long long
#define getchar() (p1==p2&&(p2=(p1=buf)+fread(buf,1,1<<21,stdin),p1==p2)?EOF:*p1++)
using namespace std;
const int mod = 998244353;
ll n,p;int k;
ll qp(ll x,int y = mod-2)
{
    ll ans = 1;
    for(;y;y >>= 1,x = x*x%mod)
        if(y&1)(ans *= x) %= mod;
    return ans;
}
char buf[1<<21],*p1,*p2;
inline int rd()
{
    char c;int f = 1;
    while(!isdigit(c = getchar()))if(c=='-')f = -1;
    int x = c-'0';
    while(isdigit(c = getchar()))x = x*10+(c^48);
    return x*f;
}
int main()
{
    // freopen(".in","r",stdin);
    // freopen(".out","w",stdout);
    n = rd();p = rd();k = rd();
    ll wn = qp(3,(mod-1)/k),w = wn,ans = (n*p%mod*qp(p+1,n-1)+qp(p+1,n))%mod;
    for(int i = 1;i < k;i++,w = wn*w%mod)
        (ans += (w*qp((p*w+1)%mod,n)-qp(p+1,n)+mod)%mod*qp(w-1)) %= mod;
    cout << (ans*qp(k)-qp(p+1,n)+mod)%mod << endl;
    return 0;
}