「NOTE」單位根反演

聽課的時候有一道題要單位根反演；
看起來是個輕量級的算法，而後就來學一下

---

# 引理

定義 $\omega_n$ 表示 $n$ 次單位復根，即 $\omega_n^n=1$，對於任意正整數 $k$，有 $$n\cdot[n\mid k]=\sum_{i=0}^{n-1}\omega_n^{ik}$$

證實只需用到單位復根的性質。

若是 \(n\mid k\)，則 \(\omega_n^{ik}=1\)，

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

此時等比數列的公比不爲 \(\mathbf1\)（\(n\mid k\) 的時候公比爲 \(1\)，因此不能用等比數列求和），直接用等比數列求和獲得

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

又由於 \(\omega_n^{nk}=(\omega_n^n)^k=1\)，因此

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

---

# 直接推導

求 \(n\) 次多項式 \(f(x)\) 的全部 \(x^{ik}\) 項的係數之和（\(k\le n\)）。形式化的，求

\[R=\sum_{i=0}^{\lfloor\tfrac ni\rfloor}[x^{ik}]f(x) \]

對式子化簡：

\[\begin{aligned} R&=\sum_{i=0}^{n-1}[k\mid i][x^i]f(x)\\ &=\sum_{i=0}^{n-1}\frac 1k\Big(\sum_{j=0}^{k-1}\omega_k^{ij}\Big)[x^i]f(x)\\ &=\sum_{i=0}^{n-1}\frac1k\sum_{j=0}^{k-1}\omega_{k}^{ij}\cdot[x^i]f(x)\\ &=\frac1k\sum_{j=0}^{k-1}\sum_{i=0}^{n-1}[x^i]f(x)\cdot(\omega_k^j)^i\\ &=\frac1k\sum_{j=0}^{k-1}f(\omega_k^j) \end{aligned} \]

就把 \(k\) 個單位復根代入式子計算。固然通常來講式子的次數很是高，可是若是式子有比較快速的計算一次的方法，單位根反演就比較有用了。

# DFT推導

另外，還能夠用 DFT 理解單位根反演，設 \(a_i\) 爲 \([x^i]f(x)\)：

\[\begin{bmatrix} \omega_k^{0\times 0}&\omega_k^{0\times 1}&\cdots&\omega_k^{0\times (n-1)}\\ \omega_k^{1\times0}&\omega_k^{1\times1}&\cdots&\omega_k^{1\times (n-1)}\\ \vdots&\vdots&\ddots&\vdots\\ \omega_k^{(k-1)\times0}&\omega_k^{(k-1)\times1}&\cdots&\omega_k^{(k-1)\times (n-1)} \end{bmatrix} \times \begin{bmatrix} a_0\\a_1\\\vdots\\a_{n-1} \end{bmatrix} = \begin{bmatrix} f(\omega_k^{0})\\f(\omega_k^{1})\\\vdots\\f(\omega_k^{k-1}) \end{bmatrix} \]

由於單位根 \(k\) 次一循環，因此能夠找到下圖所示的規律：

因而能夠將矩陣「壓縮」成 \(k\times k\) 的：

設 \(b_t=\sum a_{ik+t}\)，則壓縮後的式子就是

\[\begin{bmatrix} \omega_k^{0\times 0}&\omega_k^{0\times 1}&\cdots&\omega_k^{0\times (k-1)}\\ \omega_k^{1\times0}&\omega_k^{1\times1}&\cdots&\omega_k^{1\times (k-1)}\\ \vdots&\vdots&\ddots&\vdots\\ \omega_k^{(k-1)\times0}&\omega_k^{(k-1)\times1}&\cdots&\omega_k^{(k-1)\times (k-1)} \end{bmatrix} \times \begin{bmatrix} b_0\\b_1\\\vdots\\b_{k-1} \end{bmatrix} = \begin{bmatrix} f(\omega_k^{0})\\f(\omega_k^{1})\\\vdots\\f(\omega_k^{k-1}) \end{bmatrix} \]

不難發現上式就至關因而對 \(b_i\) 的矩陣求了一個 DFT，那麼咱們只須要 IDFT 就能夠獲得 \(b_i\) 了。

\[b_i=\frac1k\sum_{t=0}^{k-1}\omega_k^{-it}f(\omega_k^t) \]

而直接推導求出的 \(\sum a_{ik}\) 的公式則是 IDFT 的一個特例。

---

# 例題 - 生成樹計數增強版(LOJ)

一道不這麼板子題的例題

看到不進位加法，並且仍是三進制的……只能想到一種作法，拆位。那麼只用考慮邊權只有 \([0,2]\) 的狀況。

對於生成樹計數，咱們有一個很熟悉的算法叫作矩陣樹定理。矩陣樹定理最基礎的應用就是求生成樹的總方案數，可是還有一個比較經常使用的技巧：把矩陣 \(M\) 中的整數換成其餘類型——多項式。

若是存在邊 \((u,v)\) 的權值爲 \(w\)，則：

M[u][u]+=x^w,M[v][v]+=x^w;
M[u][v]-=x^w,M[v][u]-=x^w;

最後作完矩陣樹定理，咱們會獲得一個次數很高的多項式 \(F(x)\)。\(F(x)\) 第 \(i\) 位的係數 \(f_i\) 表示生成樹的邊權之和爲 \(i\) 的方案數，而咱們要求的答案是

\[0\times\sum_{i}f_{3i}+1\times\sum_{i}f_{3i+1}+2\times\sum_{i}f_{3i+2} \]

那麼只須要分別求出 \(F(x)\) 的模 \(3\) 餘 \(0,1,2\) 的項的係數便可，直接套用單位根反演——作三次行列式，算出 \(F(\omega_k^0),F(\omega_k^1),F(\omega_k^2)\)，而後再 IDFT 就能夠了。

/*Lucky_Glass*/
#include<cstdio>
#include<cassert>
#include<cstring>
#include<algorithm>
using namespace std;

const int MOD=1e9+7,VARA=500000003,VARB=541031193,N=105,M=N*N;
#define cint const int &
//(a+bi)^3=1 (mod MOD)
inline int Rint(int &r){
	int b=1,c=getchar();r=0;
	while(c<'0' || '9'<c) b=c=='-'? -1:b,c=getchar();
	while('0'<=c && c<='9') r=(r<<1)+(r<<3)+(c^'0'),c=getchar();
	return r*=b;
}
inline int Mul(cint a,cint b){return 1ll*a*b%MOD;}
inline int Pow(cint a,cint b){return b? Mul(Pow(Mul(a,a),b>>1),(b&1)? a:1):1;}
inline int Add(cint a,cint b){return a+b>=MOD? a+b-MOD:a+b;}
inline int Sub(cint a,cint b){return a-b<0? a-b+MOD:a-b;}
#define square(a) Mul(a,a)
const int INV3=Pow(3,MOD-2);

struct NCOMPLEX{
	#define cnc const NCOMPLEX &
	#define oper(typ) inline friend typ operator
	int a,b;
	NCOMPLEX(){}
	NCOMPLEX(cint A,cint B):a(A),b(B){}
	oper(NCOMPLEX) *(cnc A,cnc B){return NCOMPLEX(Sub(Mul(A.a,B.a),Mul(A.b,B.b)),Add(Mul(A.a,B.b),Mul(B.a,A.b)));}
	oper(NCOMPLEX) *(cnc A,cint B){return NCOMPLEX(Mul(A.a,B),Mul(A.b,B));}
	oper(NCOMPLEX) +(cnc A,cnc B){return NCOMPLEX(Add(A.a,B.a),Add(A.b,B.b));}
	oper(NCOMPLEX) -(cnc A,cnc B){return NCOMPLEX(Sub(A.a,B.a),Sub(A.b,B.b));}
	oper(NCOMPLEX) -(cnc A){return NCOMPLEX(Sub(0,A.a),Sub(0,A.b));}
	inline void operator +=(cnc A){(*this)=(*this)+A;}
	inline void operator -=(cnc A){(*this)=(*this)-A;}
	inline void operator *=(cnc A){(*this)=(*this)*A;}
	inline friend NCOMPLEX inv(cnc A){
		return NCOMPLEX(A.a,Sub(0,A.b))
			  *Pow(Add(square(A.a),square(A.b)),MOD-2);
	}
	#undef cnc
	#undef oper
};

const NCOMPLEX con[3]={NCOMPLEX(1,0),NCOMPLEX(VARA,VARB),NCOMPLEX(VARA,VARB)*NCOMPLEX(VARA,VARB)};

int n,m;
int edg[M][3];
NCOMPLEX mat[N][N];

NCOMPLEX Det(){
	NCOMPLEX ans(1,0);
	for(int i=1;i<n;i++){
		for(int j=i;j<n;j++)
			if(mat[j][i].a || mat[j][i].b){
				if(i==j) break;
				swap(mat[i],mat[j]),ans=-ans;
				break;
			}
		ans=ans*mat[i][i];
		NCOMPLEX tmp=inv(mat[i][i]);
		for(int j=i+1;j<n;j++){
			NCOMPLEX now=mat[j][i]*tmp;
			for(int k=i;k<n;k++) mat[j][k]-=now*mat[i][k];
		}
	}
	return ans;
}
inline void IDFT(NCOMPLEX *p){
	int i;
	NCOMPLEX tmp[3];
	tmp[0]=p[0],tmp[1]=p[1],tmp[2]=p[2];
	p[0]=tmp[0]+tmp[1]+tmp[2];
	p[1]=tmp[0]+tmp[1]*inv(con[1])+tmp[2]*inv(con[2]);
	p[2]=tmp[0]+tmp[1]*inv(con[2])+tmp[2]*inv(con[1]);
	for(i=0;i<3;++i) p[i]=p[i]*INV3;
}
int Calc(){
	NCOMPLEX res[3];
	for(int i=0;i<3;i++){
		memset(mat,0,sizeof mat);
		NCOMPLEX w[3]={con[0],con[i],con[i*2%3]};
		for(int j=1;j<=m;j++){
			int typ=edg[j][2]%3;
			mat[edg[j][0]][edg[j][0]]+=w[typ];
			mat[edg[j][1]][edg[j][1]]+=w[typ];
			mat[edg[j][0]][edg[j][1]]-=w[typ];
			mat[edg[j][1]][edg[j][0]]-=w[typ];
		}
		res[i]=Det();
	}
	IDFT(res);
	return Add(res[1].a,Add(res[2].a,res[2].a));
}
int main(){
	freopen("sum.in","r",stdin);
	freopen("sum.out","w",stdout);
	Rint(n),Rint(m);
	for(int i=1;i<=m;i++)
		Rint(edg[i][0]),Rint(edg[i][1]),Rint(edg[i][2]);
	int ans=0;
	for(int tim=1;;tim=Mul(tim,3)){
		ans=Add(ans,Mul(tim,Calc()));
		bool exi=false;
		for(int i=1;i<=m;i++)
			if(edg[i][2]/=3)
				exi=true;
		if(!exi) break;
	}
	printf("%d\n",ans);
	return 0;
}

---

THE END

Thanks for reading!

\[\begin{split} 「\ &擡頭看着星星在唱歌\\ &她的呼吸好似對我說\\ &她說你要慢慢長大\\ &不僅爲本身活着\\ &若是你也聽見星星的歌\\ &不要哭泣不要再受折磨\\ &若你擡起頭\\ &她就在天空\ 」\\ ——&\text{《星星在唱歌》By 司南} \end{split} \]

> Linked 星星在唱歌-網易雲