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;
}