【瞎口胡】单位根反演

单位根反演是用单位根来表示整除关系的东西。

定义式

\[\left[k\mid n\right] = \dfrac{1}{k} \sum \limits_{i=0}^{k-1} \omega_k^{in} \]

如果 \(k \mid n\)，那么 \(\omega_k^{n}=(\omega_k^{k})^{\frac nk}=1\)，从而 \(\omega_{k}^{in}=1\)，于是右边和式的值等于 \(k\)，乘上 \(\dfrac{1}{k}\) 后等于 \(1\)。

如果 \(k \nmid n\)，那么 \(\omega_{k}^n \neq 1\)，此时使用等比数列求和公式求得右边等于 \(\dfrac{\omega_{k}^{nk}-\omega_k^{0}}{\omega_{k}^n-1} = 0\)。

这样和式里面需要枚举整除关系的时候我们就可以用这个来代替，方便简化式子。

例题 Luogu P5591 小猪佩奇学数学

题意

给定 \(n,p,k\)，求

\[\sum_{i=0}^n \binom n i \times p^{i} \times \left\lfloor \frac{i}{k} \right\rfloor \bmod 998244353 \]

\(1 \leq n,p <998244353,k \in \{2^{w}|0 \leq w \leq 20\}\)

题解

注意到这东西不能整除分块，于是我们先对那个下取整开刀，变成 \(\dfrac{i - i \bmod k}{k}\)，然后提出 \(\dfrac{1}{k}\)，于是原式等于

\[\dfrac{1}{k} \sum \limits_{i=0}^n \dbinom{n}{i}\times p^i(i - i \bmod k) \\ = \dfrac{1}{k} \left ( \sum \limits_{i=0}^n i \times \dbinom{n}{i} p^i - \sum \limits_{i=0}^n \dbinom{n}{i} p^i \times i \bmod k\right) \]

不管外面那个 \(\dfrac{1}{k}\)，先看左边的那个和式，吸收公式

\[\sum \limits_{i=0}^n i \times \dbinom{n}{i} p^i = \sum \limits_{i=0}^n n\dbinom{n-1}{i-1} p^i \\ = n\sum \limits_{i=0}^{n-1} \dbinom{n-1}{i} p^{i+1} \\ = np \sum \limits_{i=0}^{n-1} \dbinom{n-1}{i} p^i \\ = np (1+p)^{n-1} \]

考虑右边的 \(\sum \limits_{i=0}^n \dbinom{n}{i} p^i \times i \bmod k\)。我们并没有什么办法来处理 \(i \bmod k\)，于是枚举它，

\[\sum \limits_{i=0}^n \dbinom{n}{i} p^i \times i \bmod k = \sum \limits_{r=0}^{k-1} r \sum \limits_{i=0}^{n} \dbinom{n}{i} p^i [k \mid (i - r)] \\ \]

接上单位根反演

\[\sum \limits_{r=0}^{k-1} r \sum \limits_{i=0}^{n} \dbinom{n}{i} p^i \left( \dfrac{1}{k} \sum \limits_{j=0}^{k-1} \omega_k^{j(i-r)}\right) \\ =\dfrac{1}{k} \sum \limits_{r=0}^{k-1} r \sum \limits_{i=0}^{n} \dbinom{n}{i} p^i \sum \limits_{j=0}^{k-1} \omega_k^{ij}·\omega_{k}^{-jr} \]

第二个求和号很像二项式，于是把它往里面放

\[= \dfrac 1k \sum \limits_{r=0}^{k-1} r \sum \limits_{j=0}^{k-1} \omega_{k}^{-jr} \sum \limits_{i=0}^n \dbinom{n}{i} p^i \omega_{k}^{ij} \\ = \dfrac 1k \sum \limits_{r=0}^{k-1} r \sum \limits_{j=0}^{k-1} \omega_{k}^{-jr} \sum \limits_{i=0}^n \dbinom{n}{i} (p\cdot \omega_k^j)^i \\ = \dfrac 1k \sum \limits_{r=0}^{k-1} r \sum \limits_{j=0}^{k-1} \omega_k^{-jr} (1+p\cdot \omega_k^j)^n \]

里面的一坨只和 \(j\) 一个变量相关，考虑交换求和号

\[= \dfrac 1k \sum \limits_{j=0}^{k-1} (1+p\cdot\omega_k^j)^n \sum \limits_{r=0}^{k-1} r (\omega_{k}^{-j})^r \]

第二个求和号是一个差比数列求和，求和方法和等比数列类似，考虑多算一项，和当前项做差后比较系数。形如

\[\sum \limits_{i=0}^{n-1} i\cdot a^i \]

的数列，记其和为 \(S\)，那么

\[S = \sum \limits_{i=0}^{n-1} (i+1)\cdot a^{i+1} - n\cdot a^n \\ = a \sum \limits_{i=0}^{n-1} i \cdot a^i - a \sum \limits_{i=0}^{n-1} a^i - n\cdot a^n \\ = aS - a\dfrac{a^n-1}{a-1}-n\cdot a^n \]

移项，得

\[(a-1)S = n\cdot a^n - \dfrac{a^{n+1}-a}{a-1} \\ S = \dfrac{n \cdot a^n}{a-1} - \dfrac{a^{n+1}-a}{(a-1)^2} \ \ (a \neq 1) \]

如果 \(a=1\)，那么 \(S=\dfrac{n(n-1)}{2}\)。

再来看上面的式子

\[= \dfrac 1k \sum \limits_{j=0}^{k-1} (1+p\cdot\omega_k^j)^n \sum \limits_{r=0}^{k-1} r (\omega_{k}^{-j})^r \]

预处理单位根，枚举 \(j\) 的时候用快速幂算 \((1+p\cdot \omega_k^j)^n\) 和后面的差比数列和即可，只有一个 \(\log\)。

# include <bits/stdc++.h>

const int N=1048580,INF=0x3f3f3f3f,MOD=998244353;

int w[N];
int n,p,k;

inline int read(void){
	int res,f=1;
	char c;
	while((c=getchar())<'0'||c>'9')
		if(c=='-')f=-1;
	res=c-48;
	while((c=getchar())>='0'&&c<='9')
		res=res*10+c-48;
	return res*f;
}
inline int qpow(int d,int p){
	int ans=1;
	for(;p;p>>=1,d=1ll*d*d%MOD) (p&1)&&(ans=1ll*ans*d%MOD);
	return ans;
}
inline int add(int x,int y){
	return (x+y>=MOD)?(x+y-MOD):(x+y);
}
inline int mi(int x,int y){
	return (x-y<0)?(x-y+MOD):(x-y);
}
inline int S(int r,int n){
	if(r==1) return 1ll*n*(n-1)/2%MOD;
	return mi(1ll*n*qpow(r,n)%MOD*qpow(r-1,MOD-2)%MOD,1ll*mi(qpow(r,n+1),r)*qpow(qpow(r-1,2),MOD-2)%MOD);
}

int main(void){
	n=read(),p=read(),k=read();
	int ans=1ll*n*p%MOD*qpow(p+1,n-1)%MOD,sum=0;
	w[0]=1,w[1]=qpow(3,(MOD-1)/k);
	for(int i=2;i<=k;++i) w[i]=1ll*w[i-1]*w[1]%MOD;
	for(int i=0;i<k;++i) sum=add(sum,1ll*qpow(add(1ll*w[i]*p%MOD,1),n)*S(w[k-i],k)%MOD);
	ans=mi(ans,1ll*sum*qpow(k,MOD-2)%MOD);
	printf("%lld",1ll*ans*qpow(k,MOD-2)%MOD);
	return 0;
}