单位根反演小记

反演公式

\[[n | v] = \frac{1}{n}\sum_{0 \le j < n}(\omega_n^v)^j \]

证明很简单，等比数列求和即可。

计算单位根

一般来说我们都是计算原根，而对于一个素数 \(p\)，\(g\) 是原根当且仅当对于所有 \(p-1\) 的质因子 \(p_i\) 都有 \(g^{\frac{p-1}{p_i}} \not\equiv 1 \pmod p\)。

而最小质因子是 \(O(m^{\frac{1}{4}})\)，所以我们暴力枚举加检验即可。

int cal() {//找到 mod 的最小原根 
	int phi = mod - 1, cur = 0;//分解质因数
	for (int d = 2; d * d <= phi; d++) {
		if (phi % d == 0) {
			p[++cur] = d;
			while (phi % d == 0)
				phi /= d;
		}
	}
	if (phi > 1)
		p[++cur] = phi;
	int g = 0;
	for (int i = 1; i < mod; i++) {
		bool flg = true;
		for (int j = 1; j <= cur; j++)
			if (fpow(i, (mod - 1) / p[j], mod) == 1) {
				flg = false;
				break;
			}
		if (flg) {
			g = i;
			break;
		}
	}
	return g;
}

应用

一定要注意 \(\omega_{n}^0 = 1\)！！！！

牛客Wannafly挑战赛11E 白兔的***难

题意：

给定 \(k \le 2^20, n \le 10^{16}, p = 998244353\)，求 \(t \in [0, k)\)，\(a_t =\sum_{k | i,0 \le i + t \le n} \binom{n}{i + t}\)。

思路：

直接单位根反演：

\[\begin{aligned} \sum_{k | i,0 \le i + t \le n} \binom{n}{i + t} &= \sum_{0 \le i\le n} \binom{n}{i}[k | i - t]\\ &= \sum_{0 \le i\le n} \binom{n}{i}\frac{1}{k}\sum_{j = 0}^{k-1}(\omega_{k}^{i -t})^j\\ &= \frac{1}{k}\sum_{j = 0}^{k-1}\omega_{k}^{-tj}\sum_{0 \le i\le n} \binom{n}{i}\omega_{k}^{ij}\\ &= \frac{1}{k}\sum_{j = 0}^{k-1}\omega_{k}^{-tj}(1 + \omega_k^j)^n\\ \end{aligned} \]

不妨令 \(y_j = (1 + \omega_k^j)^n\)，则上式就是一个 IDFT 形式的东西，利用 NTT 解决即可。

P5293 [HNOI2019] 白兔之舞

设 \(W\) 是 \(n \times n\) 的邻接矩阵，则我们有走 \(m\) 步的方法数等于：

\[\binom{L}{m}W^m \]

所以我们要求的就是对于所有 \(t\)：

\[\sum_{m=0}^L\binom{L}{m}W^m[k | m - t] \]

单位根反演可以化为：

\[\frac{1}{k}\sum_{j=0}^{k-1}\omega_k^{-tj}(1 + \omega_k^jW)^m \]

不妨设 \(A_j = (I + \omega_{k}^jW)^L\)，\(a_j\) 为 \(W\) 中编号第 \(j\) 个的系数，则我们要求的就是 \(IDFT(a)\)。使用 MTT 即可解决。

【牛客Wannafly挑战赛23 F】计数

直接大力推式子：

\[\begin{aligned} ANS &= \sum_{T}[k | \sum_{e \in T}val_e]\\ &= \frac{1}{k}\sum_{T}\sum_{j=0}^{k-1}(\omega_{k}^{\sum_{e \in T}val_e})^j\\ &= \frac{1}{k}\sum_{j=0}^{k-1}\sum_{T}\prod_{e \in T}(\omega_k^{j})^{val_e} \end{aligned} \]

转化成标准的矩阵树定理形式求解，时间复杂度 \(O(kn^3)\)。

LOJ6485 LJJ 学二项式定理

考虑对余数分类讨论，假设 \(i \bmod 4 = k\)，则贡献为：

\[a_k\sum_{i=0}^n\binom{n}{i}s^{i}[4 | i - k] \]

我们还是利用单位根反演，最终可以推出：

\[\frac{1}{4}a_k\sum_{j=0}^3\omega_{4}^{-kj}(1 + s\omega_4^j)^n \]

近似于 \(O(1)\) 计算了。

P5591 小猪佩奇学数学

很好的一道题。

首先这道题明显复杂度和 \(k\) 相关，且 \(k\) 的特性以及模数是 \(998244353\) 告诉我们这道题很明显要让我们使用单位根来做。

我们希望把式子化成存在 \(\bmod\) 的形式，不妨考虑下面这个式子：

\[\lfloor\frac{i}{k}\rfloor = \frac{i - i \bmod k}{k} \]

所以我们将计算的式子变为：

\[\sum_{i=0}^{n}\binom{n}{i}p^i\frac{i - i \bmod k}{k} \]

我们分别计算两部分，对于：

\[\sum_{i=0}^ni\binom{n}{i}p^i \]

我们根据组合数的重要性质 \(\binom{n}{i} = \frac{n}{i}\binom{n - 1}{i - 1}\)，得到 \(i\binom{n}{i} = n\binom{n-1}{i-1}\)，所以我们可以将其转化为标准的二项式定理，从而得到其封闭形式：

\[np(p + 1)^{n-1} \]

对于后面的式子，枚举余数 \(r\)，可以得到：

\[a_r = \frac{1}{k}\sum_{j=0}^{k-1}\omega_{k}^{-rj}(p\omega_{k}^j+1)^n \]

直接 IDFT 即可，时间复杂度 \(O(k \log k)\)。

Guideposts OpenJ_POJ - 1058

类似白兔之舞可以转化成：

\[\frac{1}{k}\sum_{j = 0}^{k - 1}(I + \omega_k^jW)^n \]

于是我们每次只需要计算矩阵快速幂，然后将 \((s, t)\) 作为答案加起来即可。

BZOJ3328 PYXFIB

转化成：

\[\sum_{i=0}^n\binom{n}{i}F_i[k|i] \]

然后单位根反演，用矩阵快速幂来计算即可。

UOJ450 复读机

绝世好题，被我推出来了，赢（）

先考虑答案，我们可以写一个暴力一点的式子：

\[\sum_{i_1 + i_2 + \dots + i_k = n}\binom{n}{i_1, i_2, \dots, i_k}[d|i_1][d|i_2]\dots[d|i_k] \]

这就是答案。

考虑计算，对于这种带有多项式系数的式子不难想到 EGF，我们定义：

\[F(x) = \sum_{i}\frac{[d|i]x^i}{i!} \]

则我们要求的就是 \(n![x^n]F^k(x)\)，考虑对于 \(F(x)\) 单位根反演：

\[\begin{aligned} F(x) &= \sum_{i}\frac{[d|i]x^i}{i!}\\ &= \sum_{i}\frac{x^i}{i!}\frac{1}{d}\sum_{j=0}^{d-1}(\omega_d^i)^j\\ &= \frac{1}{d}\sum_{j=0}^{d-1}\sum_{i}\frac{(x\omega_d^j)^i}{i!}\\ &= \frac{1}{d}\sum_{j=0}^{d-1}e^{x\omega_d^j}\\ \end{aligned} \]

则我们要求的就是：

\[\frac{1}{d^k}[x^n](\sum_{j=0}^{d-1}e^{x\omega_d^j}) \]

发现 \(d\) 只有三种取值，我们分类讨论：

对于 \(d=1\) 答案为：

\[\frac{n!}{d^k}[x^n]e^{xk} \]

直接 taylor 展开得到：

\[\frac{n!}{d^k}\frac{k^n}{n!} = k^n \]

考虑组合意义不难发现相当于没有任何限制，每一秒可以任意选一个复读机。

对于 \(d=2\)，答案为：

\[\begin{aligned} \frac{n!}{d^k}[x^n](e^x + e^{x\omega})^k &= \frac{n!}{d^k}\sum_{a+b=k}\binom{k}{a}[x^n]e^{ax + bx\omega}\\ &= \frac{n!}{d^k}\sum_{a+b=k}\binom{k}{a}[x^n]e^{x(a + b\omega)}\\ &= \frac{n!}{d^k}\sum_{a+b=k}\binom{k}{a}[x^n]\sum_{j}\frac{x^j(a+b\omega)^j}{j!}\\ &= \frac{n!}{d^k}\sum_{a+b=k}\binom{k}{a}\frac{(a+b\omega)^n}{n!}\\ &= \frac{1}{d^k}\sum_{a+b=k}\binom{k}{a}(a+b\omega)^n\\ \end{aligned} \]

\(O(k)\) 枚举计算即可。

对于 \(d=3\)，类似的，我们可以得到最后答案为：

\[\frac{1}{d^k}\sum_{a+b+c=k}\binom{k}{a,b,c}(a+b\omega+c\omega^2)^n \]

可以 \(O(k^2)\) 计算。这道题就做完了。

个人觉得是单位根反演题中不那么套路的题。