P5591-小猪佩奇学数学【单位根反演】

正题

题目链接:https://www.luogu.com.cn/problem/P5591

---

题目大意

给出(n,p,k)求

[left(sum_{i=0}^ninom{n}{i}p^ileftlfloorfrac{i}{k} ight floor ight)mod 998244353 ]

(1leq n,p<998244353,k=2^w(win[0,20]))

---

解题思路

开始以为推错了，结果是要特判
推出了看上去不是我能推的式子

[sum_{i=1}^ninom{n}{i}p^isum_{j=1}^i[k|j] ]

然后单位根反演

[sum_{i=1}^ninom{n}{i}p^isum_{j=1}^ifrac{1}{k}sum_{l=0}^{k-1}omega_k^{l imes j} ]

系统整理一下

[frac{1}{k}sum_{l=0}^{k-1}sum_{i=1}^ninom{n}{i}p^isum_{j=1}^iomega_k^{l imes j} ]

然后等比数列通项公式拆开

[frac{1}{k}sum_{l=0}^{k-1}sum_{i=1}^ninom{n}{i}p^ifrac{omega_k^l-omega_{k}^{li}omega^l_k}{1-omega_k^l} ]

[frac{1}{k}sum_{l=0}^kfrac{omega_k^l}{1-omega_k^l}left(sum_{i=1}^ninom{n}{i}p^i-sum_{i=1}^ninom{n}{i}p^iomega_k^{li} ight) ]

[frac{1}{k}sum_{l=0}^kfrac{omega_k^l}{1-omega_k^l}left((p+1)^n-(pomega_k^l+1)^n ight) ]

然后写出来会愉快的发现没有过样例，仔细看我们式子里面有一个(frac{omega_k^l}{1-omega_k^l})。

当 (l=0) 的时候(1-omega_k^l=0)，所以不能直接这么求，我们这个得分开考虑。

就是

[sum_{i=1}^ninom{n}{i}p^iiRightarrow sum_{i=1}^nfrac{n!}{(i-1)!(n-i)!}p^i ]

[npsum_{i=1}^nfrac{(n-1)!}{(i-1)!(n-i)!}p^{i-1}Rightarrow nsum_{i=1}^ninom{n-1}{i}p^{i-1}Rightarrow np(p+1)^{n-1} ]

就好了

时间复杂度(:O(klog 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
",ans*power(k,P-2)%P);
	return 0;
}