单位根反演

处理整除时才有贡献的利器。

1. 代换公式

前置知识：FFT / NTT。

如果你对 FFT 的证明过程比较熟悉的话，你应该会记得一个叫单位根的东西 $\omega_{n}^k$，他还有一个特殊的性质：

$$\dfrac 1n \sum_{i = 0} ^ {n - 1} \omega_{n}^{ik} = [n | k]$$

如何证明这个结论呢？

当 $k\bmod n\not= 0$ 时，我们的公比为 $\omega_{n}^k$ 不是 1，考虑使用等比数列求和公式：

$$\sum_{i = 0} ^ {n - 1} \omega_{n}^{ik} = \dfrac{\omega_n^{nk} - 1}{\omega_{n}^k - 1}$$

上面的 $\omega_{n}^{nk}$ 等于 1，所以整个式子等于 0。原式成立。

当 $k\bmod n = 0$ 时，每一项都是 1，求和后除以 $n$ 就是 1 了。

这个有什么用了？当遇到 $[i\bmod n = j]$ 这种情况时，我们可以考虑时候单位根反演展开，和其他项合并以快速计算。我们来看几道例题。

2. 例题

T1：LJJ 学二项式定理

题目传送门 LOJ

看到 $a_{i\bmod 4}$，这个就是经典的单位根反演了。

首先我们化成上面的那个形式：

$$ans = \sum_{i = 0} ^ n \binom ni s ^ i \sum_{j = 0} ^ 3 a_j [i\bmod 4 = j]$$

$[i\bmod 4 = j] = [4 | (i - j)]$，然后对后面这个式子用单位根反演大力展开：

$$\sum_{j = 0} ^ 3 a_j[i\bmod 4 = j] = \sum_{j = 0} ^ 3 \dfrac{a_j}{4} \sum_{k = 0} ^ 3 \omega_{4}^{ik - jk}$$

然后 $n$ 在前面是没有出路的，我们考虑将 $j, k$ 放在前面，分离单位根的 $i, j$ 变量：

$$ans = \sum_{k = 0} ^ 3 \sum_{j = 0} ^ 3 \omega_{4}^{-jk}\dfrac{a_j}4 \sum_{i = 0} ^ n \binom ni s ^ i \omega_{4}^{ik}$$

后面是一个二项式定理，我们直接合并为 $(s \omega_{4}^k + 1) ^ n$。这样暴力枚举 $j, k$，快速幂即可做到 $O(\log n)$ 单次。处理单位根按照 NTT 里面的求法即可。

void init()
{
	wn[0] = 1, wn[1] = qpow(3, (Mod - 1) >> 2);
	wn[2] = Mod - 1, wn[3] = qpow(3, (Mod - 1) / 4 * 3);
	inv4 = qpow(4);
}

void work()
{
	int a[4], s, res = 0;
	LL n;
	scanf("%lld %d %d %d %d %d", &n, &s, a, a + 1, a + 2, a + 3);
	n %= (Mod - 1);
	for (int i = 0; i < 4; ++ i)
		for (int j = 0; j < 4; ++ j)
			res = (res + (LL) a[i] * wn[(4 - (i * j & 3)) & 3] % Mod
				* qpow((LL) s * wn[j] % Mod + 1, n)) % Mod;
	res = (LL) res * inv4 % Mod;
	printf("%d\n", res);
}

T2：PYXFIB

题目传送门 HydroOJ - BZOJ

答案容易写作：

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

其中 $F_i$ 表示第 $i$ 项斐波那契数。

还是直接单位根反演：

$$\begin{aligned} ans =& \sum_{i = 0} ^ n \binom ni \dfrac{F_i}{k} \sum_{j = 0} ^ {k - 1} \omega_k^{ij}\\ =& \dfrac 1k \sum_{j = 0} ^ {k - 1} \sum_{i = 0} ^ n \binom ni F_i \omega_k^{ij} \end{aligned}$$

现在我们有一个 $F_i$ 无法化开，不能像上面一样变形。但是我们知道 $F_n = \dfrac 1{\sqrt 5}\left((\dfrac {1 + \sqrt 5}2) ^ n - (\dfrac{1 - \sqrt 5}2) ^ n \right)$，如果 5 在该质数下有二次剩余，我们可以写作幂次的加减。但很可惜，这道题没有给出这样的条件。

看到斐波那契数，容易想到矩阵乘法，注意还是需要把次幂联系起来，否则无法计算。

直接将转移矩阵的 $n$ 次方写入，那么 $F_i = mat ^ i_{1, 1}$。于是我们直接带入可得：

$$ans = \dfrac 1k \sum_{j = 0} ^ {k - 1} \sum_{i = 0} ^ n \binom ni \omega_k^{ij}mat ^ i_{1, 1}$$

考虑前面我们处理 $\sum_{i = 0} ^ n \binom ni x ^ i$ 的时候，我们相当于是添上了乘法单位元 1 的 $n - i$ 次方。这里同样，我们添上单位矩阵 $I$ 的 $n - i$ 次方，于是答案可以写作：

$$ans = \dfrac 1k\sum_{j = 0} ^ {k - 1} (\omega_k^{j}mat + I) ^ n_{1, 1}$$

直接计算即可，复杂度 $O(k\log n)$。

int findrt()
{
	std::vector<int> fac;
	int cur = Mod - 1;
	for (int i = 2; i <= cur / i; ++ i)
		if (cur % i == 0) {
			fac.push_back(i);
			while (cur % i == 0) cur /= i;
		}
	if (cur ^ 1) fac.push_back(cur);
	auto check = [&](int x) {
		for (int t : fac)
			if (qpow(x, (Mod - 1) / t) == 1) return false;
		return true;
	};
	for (int i = 1; i < Mod; ++ i)
		if (check(i)) return i;
	return -1;
}

void work()
{
	LL n;
	int k, res = 0;
	std::cin >> n >> k >> Mod;
	wn[0] = 1, wn[1] = qpow(findrt(), (Mod - 1) / k);
	for (int i = 2; i < k; ++ i) wn[i] = (LL) wn[i - 1] * wn[1] % Mod;
	for (int i = 0; i < k; ++ i)
	{
		Matrix cur{{{1, wn[i]}, {wn[i], wn[i] + 1}}};
		cur = qpow(cur, n), adj(res += cur[1][1] - Mod);
	}
	res = (LL) res * qpow(k) % Mod;
	std::cout << res << std::endl;
}

T3：小猪佩奇学数学

题目传送门 Luogu

保证了 $k$ 是 2 的次幂，显然就需要单位根了，否则单位根就不好求了……

看到 $\lfloor\dfrac ik \rfloor$ 很神秘，先乘 $k$ 最后除回去，那么 $k\lfloor \dfrac ik\rfloor = i - i\bmod k$。

看到 $i\bmod k$，果断化为 $\sum_{j = 0} ^ {k - 1} j[k | (i - j)]$，然后单位根反演，于是就可以得到：

$$i\bmod k = \sum_{j = 0} ^ {k - 1} \dfrac jk\sum_{l = 0} ^ {k - 1} \omega_k ^ {il - jl}$$

带回原式：

$$k\times ans = \sum_{i = 0} ^ n \binom ni p ^ i (i - \sum_{j = 0} ^ {k - 1} \dfrac jk \sum_{j = 0} ^ {k - 1}\omega_k^{il - jl})$$

原式化为了两边，分别计算。

$$ans1 = \sum_{i = 0} ^ n\binom ni p ^ i i$$

孤零零的 $i$ 很烦，我们考虑将其合并进入组合数：

$$\begin{aligned} ans1 =& \sum_{i = 0} ^ n \dfrac{n!p ^ ii}{i!(n - i)!}\\ =& n\sum_{i = 1} ^ n \dfrac{(n - 1)!p ^ i}{(i - 1)!(n - i)!}\\ =& np\sum_{i = 1} ^ n \binom{n - 1}{i - 1} p ^ {i - 1}\\ =& np (p + 1) ^ {n - 1} \end{aligned}$$

这样第一部分可以在 $O(\log n)$ 的时间内解决。

接下来考虑第二部分：

$$\begin{aligned} k\times ans2 =& \sum_{i = 0} ^ n \binom ni p ^ i\sum_{j = 0} ^ {k - 1} j\sum_{l = 0} ^ {k - 1} \omega_k^{ik - jk} \end{aligned}$$

仍然考虑先枚举 $j, l$，那么可以写作：

$$\begin{aligned} k\times ans2 =& \sum_{j = 0} ^ {k - 1}j \sum_{l = 0} ^ {k - 1} \omega_k^{-jl} \sum_{i = 0} ^ {n}\binom ni p ^ i \omega_k^{il}\\ =& \sum_{j = 0} ^ {k - 1}j \sum_{l = 0} ^ {k - 1} \omega_k^{-jl}(p\omega_k^{l} + 1) ^ n \end{aligned}$$

这样就得到了 $O(k ^ 2\log n)$ 的做法，可以得 60 pts。

后面 $l$ 一坨不好化开，我们看到 $j$ 相关的形式比较少，于是交换 $j, l$ 顺序：

$$k\times ans2 = \sum_{l = 0} ^ {k - 1} (p\omega_k ^{l} + 1) ^ n\sum_{j = 0} ^ {k - 1} j\omega_k^{-jl}$$

这里后面一坨还是不好计算，但是由于单位根特殊的性质 $\omega_k^{nk} = 0$，我们可以考虑错位减一下。

设 $S(x) = \sum_{j = 0} ^ {k - 1} j\omega_k^{jx}$，则有：

$$\begin{aligned} S(x) =& \sum_{j = 0} ^ {k - 1} j\omega_k^{jx}\\ \omega_k^{x} S(x) =& \sum_{j = 1} ^ k (j - 1)\omega_k^{jx}\\ (\omega_k^{x} - 1)S(x) =& (k - 1)\omega_k^{jk} -\sum_{j = 1} ^ {k - 1} \omega_k^{jx} \end{aligned}$$

前面的 $(k - 1)\omega_k^{jk}$ 就是 $k - 1$，后面的 $\sum_{j = 1} ^ {k - 1} \omega_k^{jx}$ 其实是 -1，因为我们尝试加上 $\omega_k^0$ 整个式子就变成 0 了。

那么我们可以得到 $S(x) = \dfrac{k}{\omega_k^{x} - 1}$，唯一注意 $n | x$ 的时候 $S(x) = \dfrac{n(n - 1)}2$。总复杂度 $O(k\log n)$，可以通过。

int wnsum(int n, int mul)
{
	if (mul == n) return (LL) n * (n - 1) / 2 % Mod;
	else return (LL) n * qpow(wn[mul] - 1) % Mod;
}

int main()
{
	std::cin >> n >> p >> k;
	wn[0] = 1, wn[1] = qpow(3, (Mod - 1) / k); 
	for (int i = 2; i < k; ++ i) wn[i] = (LL) wn[i - 1] * wn[1] % Mod;
	int res = (LL) n * qpow(p + 1, n - 1) % Mod * p % Mod, res2 = 0;
	for (int l = 0; l < k; ++ l)
		res2 = (res2 + (LL) qpow((LL) p * wn[l] % Mod + 1, n) * wnsum(k, k - l)) % Mod;
	res2 = res2 * (LL) qpow(k) % Mod;
	adj(res -= res2), res = (LL) res * qpow(k) % Mod;
	printf("%d\n", res);
	return 0;
}