P5591-小猪佩奇学数学【单位根反演】
正题
题目链接:https://www.luogu.com.cn/problem/P5591
题目大意
给出\(n,p,k\)求
\[\left(\sum_{i=0}^n\binom{n}{i}p^i\left\lfloor\frac{i}{k}\right\rfloor \right)\mod 998244353
\]
\(1\leq n,p<998244353,k=2^w(w\in[0,20])\)
解题思路
开始以为推错了,结果是要特判
推出了看上去不是我能推的式子
\[\sum_{i=1}^n\binom{n}{i}p^i\sum_{j=1}^i[k|j]
\]
然后单位根反演
\[\sum_{i=1}^n\binom{n}{i}p^i\sum_{j=1}^i\frac{1}{k}\sum_{l=0}^{k-1}\omega_k^{l\times j}
\]
系统整理一下
\[\frac{1}{k}\sum_{l=0}^{k-1}\sum_{i=1}^n\binom{n}{i}p^i\sum_{j=1}^i\omega_k^{l\times j}
\]
然后等比数列通项公式拆开
\[\frac{1}{k}\sum_{l=0}^{k-1}\sum_{i=1}^n\binom{n}{i}p^i\frac{\omega_k^l-\omega_{k}^{li}\omega^l_k}{1-\omega_k^l}
\]
\[\frac{1}{k}\sum_{l=0}^k\frac{\omega_k^l}{1-\omega_k^l}\left(\sum_{i=1}^n\binom{n}{i}p^i-\sum_{i=1}^n\binom{n}{i}p^i\omega_k^{li}\right)
\]
\[\frac{1}{k}\sum_{l=0}^k\frac{\omega_k^l}{1-\omega_k^l}\left((p+1)^n-(p\omega_k^l+1)^n\right)
\]
然后写出来会愉快的发现没有过样例,仔细看我们式子里面有一个\(\frac{\omega_k^l}{1-\omega_k^l}\)。
当 \(l=0\) 的时候\(1-\omega_k^l=0\),所以不能直接这么求,我们这个得分开考虑。
就是
\[\sum_{i=1}^n\binom{n}{i}p^ii\Rightarrow \sum_{i=1}^n\frac{n!}{(i-1)!(n-i)!}p^i
\]
\[np\sum_{i=1}^n\frac{(n-1)!}{(i-1)!(n-i)!}p^{i-1}\Rightarrow n\sum_{i=1}^n\binom{n-1}{i}p^{i-1}\Rightarrow np(p+1)^{n-1}
\]
就好了
时间复杂度\(:O(k\log P)\)
code
#include<cstdio>
#include<cstring>
#include<algorithm>
#define ll long long
using namespace std;
const ll P=998244353;
ll n,p,k,ans;
ll power(ll x,ll b){
ll ans=1;x%=P;
while(b){
if(b&1)ans=ans*x%P;
x=x*x%P;b>>=1;
}
return ans;
}
signed main()
{
scanf("%lld%lld%lld",&n,&p,&k);
ll d=power(3,(P-1)/k);
ans=n*p%P*power(p+1,n-1)%P;
for(ll i=1,w=d;i<k;i++,w=w*d%P){
ll inv=power(P+1-w,P-2)*w%P;
ans+=power(p+1,n)*inv%P;
ans-=power(w*p+1,n)*inv%P;
ans=(ans+P)%P;
}
printf("%lld\n",ans*power(k,P-2)%P);
return 0;
}