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

正题

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

题目大意

给出 $n, p, k$ 求
$(∑i=0n(ni)pi⌊ik⌋)mod998244353\left(\sum_{i=0}^n\binom{n}{i}p^i\left\lfloor\frac{i}{k}\right\rfloor \right)\mod 998244353$

$1≤n,p<998244353,k=2w(w∈[0,20])1\leq n,p<998244353,k=2^w(w\in[0,20])$

解题思路

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

$∑i=1n(ni)pi∑j=1i[k∣j]\sum_{i=1}^n\binom{n}{i}p^i\sum_{j=1}^i[k|j]$
然后单位根反演
$∑i=1n(ni)pi∑j=1i1k∑l=0k−1ωkl×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}$
系统整理一下
$1k∑l=0k−1∑i=1n(ni)pi∑j=1iωkl×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}$
然后等比数列通项公式拆开
$1k∑l=0k−1∑i=1n(ni)piωkl−ωkliωkl1−ωkl\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}$
$1k∑l=0kωkl1−ωkl(∑i=1n(ni)pi−∑i=1n(ni)piωkli)\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)$
$1k∑l=0kωkl1−ωkl((p+1)n−(pωkl+1)n)\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)$

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

当 $l = 0$ 的时候 $1−ωkl=01-\omega_k^l=0$ ，所以不能直接这么求，我们这个得分开考虑。

就是
$∑i=1n(ni)pii⇒∑i=1nn!(i−1)!(n−i)!pi\sum_{i=1}^n\binom{n}{i}p^ii\Rightarrow \sum_{i=1}^n\frac{n!}{(i-1)!(n-i)!}p^i$
$np∑i=1n(n−1)!(i−1)!(n−i)!pi−1⇒n∑i=1n(n−1i)pi−1⇒np(p+1)n−1np\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(klog⁡P):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;
}