P5591 小猪佩奇学数学

知识点

二项式定理

\[(x+1)^n=\sum_{i=0}^n\binom nix^i
\]

单位根反演

\[[n\mid k]=\frac 1n\sum_{i=0}^{n-1}\omega_n^{ik}
\]

证明：

\[[n\mid k]=\begin{cases}\frac 1n\sum_{i=0}^{n-1}\omega_n^{ik}=\frac 1n\sum_{i=0}^{n-1}1=1 &,n\mid k\\
\frac 1n\sum_{i=0}^{n-1}\omega_n^{ik}=\frac 1n\frac{\omega_n^k-\omega_n^{nk}}{1-\omega_n^k}=0 &,n\not\mid k\end{cases}
\]

题意

求

\[\sum_{i=0}^n\binom ni p^i\left\lfloor\frac ik\right\rfloor \pmod{998244353}
\]

$1\le n,p\le998244353,k\in\{2^w|0\le w\le 20\}$

思路

一看到前面这个形式容易想到二项式定理，但是后面这个 $\left\lfloor\frac ik\right\rfloor$ 不好处理。

观察一下数据范围发现 $k$ 较小，考虑使用单位根反演，我们将柿子往这边化：

\[\left\lfloor\frac ik\right\rfloor=(\sum_{j=0}^i[k\mid j])-1=(\sum_{j=0}^i[k\mid j])-1=(\sum_{j=0}^i\frac 1k\sum_{i=0}^{k-1}\omega_k^{ij})-1
\]

代入得到

\[\begin{aligned}&\sum_{i=0}^n\binom ni p^i\left\lfloor\frac ik\right\rfloor\\
&=\sum_{i=0}^n\binom ni p^i\sum_{j=0}^i\frac 1k\sum_{d=0}^{k-1}\omega_k^{dj}-(\sum_{i=0}^n\binom ni p^i)\\
&=\frac 1k\sum_{d=0}^{k-1}\sum_{i=0}^n\binom ni p^i\sum_{j=0}^i\omega_k^{dj}-(p+1)^n\\
&=\frac 1k(P+\sum_{d=1}^{k-1}\sum_{i=0}^n\binom ni p^i\frac{\omega_k^{di+d}-1}{\omega_k^d-1})-(p+1)^n\\
&=\frac 1k(P+\sum_{d=1}^{k-1}\frac{\sum_{i=0}^n\binom ni p^i(\omega_k^{di+d}-1)}{{\omega_k^d-1}})-(p+1)^n\\
&=\frac 1k(P+\sum_{d=1}^{k-1}\frac{\sum_{i=0}^n\binom ni p^i\omega_k^{di+d}-\sum_{i=0}^n\binom ni p^i}{{\omega_k^d-1}})-(p+1)^n\\
&=\frac 1k(P+\sum_{d=1}^{k-1}\frac{\omega_k^d\sum_{i=0}^n\binom ni p^i\omega_k^{di}-\sum_{i=0}^n\binom ni p^i}{{\omega_k^d-1}})-(p+1)^n\\
&=\frac 1k(P+\sum_{d=1}^{k-1}\frac{\omega_k^d(p\omega_k^d+1)^n-(p+1)^n}{{\omega_k^d-1}})-(p+1)^n
\end{aligned}
\]

上式中 $P$ 是 $d$ 等于零的情况，此时 $\sum_{j=0}^i\omega_k^{dj}$ 全为 1，公比为 1，不适用等比数列求和公式，我们单独算一下。由 $\binom nmm=\binom {n-1}{m-1}n$，有

\[\begin{aligned}
P&=\sum_{i=0}^n\binom ni p^i(i+1)\\
&=(\sum_{i=0}^n\binom ni p^ii)+(p+1)^n\\
&=(np\sum_{i=0}^n\binom {n-1}{i-1} p^{i-1})+(p+1)^n\\
&=np(p+1)^{n-1}+(p+1)^n
\end{aligned}
\]

代码

#include<iostream>

#include<cstdio>

#include<algorithm>

#include<cctype>

#include<cstring>

#include<cmath>

using namespace std;

inline int read(){

	int w=0,x=0;char c=getchar();

	while(!isdigit(c))w|=c=='-',c=getchar();

	while(isdigit(c))x=x*10+(c^48),c=getchar();

	return w?-x:x;

}

namespace star

{

	const int maxn=(1<<10)+5,mod=998244353,g=3;

	int n,p,k,ans,rt;

	inline int fpow(int a,int b){int ans=1;for(;b;b>>=1,a=1ll*a*a%mod) if(b&1) ans=1ll*ans*a%mod; return ans;}

	inline void work(){

		n=read(),p=read(),k=read(),rt=fpow(g,(mod-1)/k);

		ans=(1ll*n*p%mod*fpow(p+1,n-1)+fpow(p+1,n))%mod;

		for(int mul=rt,d=1;d<k;d++,mul=1ll*mul*rt%mod) ans=(ans+(1ll*mul*fpow((1ll*p*mul+1)%mod,n)%mod-fpow(p+1,n)+mod)%mod*fpow((mul-1+mod)%mod,mod-2))%mod;

		ans=1ll*ans*fpow(k,mod-2)%mod;

		printf("%d\n",(ans-fpow(p+1,n)+mod)%mod);

	}

}

signed main(){

	star::work();

	return 0;

}