单位根反演

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

1. 代换公式

前置知识：FFT / NTT。

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

\frac{1}{n} \sum_{i = 0}^{n - 1} ω_{n}^{i k} = [n | k]

如何证明这个结论呢？

当 $k mod n \neq 0$ 时，我们的公比为 $ω_{n}^{k}$ 不是 1，考虑使用等比数列求和公式：

\sum_{i = 0}^{n - 1} ω_{n}^{i k} = \frac{ω_{n}^{n k} - 1}{ω_{n}^{k} - 1}

上面的 $ω_{n}^{n k}$ 等于 1，所以整个式子等于 0。原式成立。

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

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

2. 例题

T1：LJJ 学二项式定理

题目传送门 LOJ

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

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

a n s = \sum_{i = 0}^{n} (\binom{n}{i}) s^{i} \sum_{j = 0}^{3} a_{j} [i mod 4 = j]

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

\sum_{j = 0}^{3} a_{j} [i mod 4 = j] = \sum_{j = 0}^{3} \frac{a_{j}}{4} \sum_{k = 0}^{3} ω_{4}^{i k - j k}

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

a n s = \sum_{k = 0}^{3} \sum_{j = 0}^{3} ω_{4}^{- j k} \frac{a_{j}}{4} \sum_{i = 0}^{n} (\binom{n}{i}) s^{i} ω_{4}^{i k}

后面是一个二项式定理，我们直接合并为 $(s ω_{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

答案容易写作：

a n s = \sum_{i = 0}^{n} [k | i] (\binom{n}{i}) F_{i}

其中 $F_{i}$ 表示第 $i$ 项斐波那契数。

还是直接单位根反演：

\begin{aligned} a n s = & \sum_{i = 0}^{n} (\binom{n}{i}) \frac{F_{i}}{k} \sum_{j = 0}^{k - 1} ω_{k}^{i j} \\ = & \frac{1}{k} \sum_{j = 0}^{k - 1} \sum_{i = 0}^{n} (\binom{n}{i}) F_{i} ω_{k}^{i j} \end{aligned}

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

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

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

a n s = \frac{1}{k} \sum_{j = 0}^{k - 1} \sum_{i = 0}^{n} (\binom{n}{i}) ω_{k}^{i j} m a t_{1, 1}^{i}

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

a n s = \frac{1}{k} \sum_{j = 0}^{k - 1} (ω_{k}^{j} m a t + I)_{1, 1}^{n}

直接计算即可，复杂度 $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 的次幂，显然就需要单位根了，否则单位根就不好求了……

看到 $⌊ \frac{i}{k} ⌋$ 很神秘，先乘 $k$ 最后除回去，那么 $k ⌊ \frac{i}{k} ⌋ = i - i mod k$ 。

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

i mod k = \sum_{j = 0}^{k - 1} \frac{j}{k} \sum_{l = 0}^{k - 1} ω_{k}^{i l - j l}

带回原式：

k \times a n s = \sum_{i = 0}^{n} (\binom{n}{i}) p^{i} (i - \sum_{j = 0}^{k - 1} \frac{j}{k} \sum_{j = 0}^{k - 1} ω_{k}^{i l - j l})

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

a n s 1 = \sum_{i = 0}^{n} (\binom{n}{i}) p^{i} i

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

\begin{aligned} a n s 1 = & \sum_{i = 0}^{n} \frac{n! p^{i} i}{i! (n - i)!} \\ = & n \sum_{i = 1}^{n} \frac{(n - 1)! p^{i}}{(i - 1)! (n - i)!} \\ = & n p \sum_{i = 1}^{n} (\binom{n - 1}{i - 1}) p^{i - 1} \\ = & n p (p + 1)^{n - 1} \end{aligned}

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

接下来考虑第二部分：

\begin{aligned} k \times a n s 2 = & \sum_{i = 0}^{n} (\binom{n}{i}) p^{i} \sum_{j = 0}^{k - 1} j \sum_{l = 0}^{k - 1} ω_{k}^{i k - j k} \end{aligned}

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

\begin{aligned} k \times a n s 2 = & \sum_{j = 0}^{k - 1} j \sum_{l = 0}^{k - 1} ω_{k}^{- j l} \sum_{i = 0}^{n} (\binom{n}{i}) p^{i} ω_{k}^{i l} \\ = & \sum_{j = 0}^{k - 1} j \sum_{l = 0}^{k - 1} ω_{k}^{- j l} (p ω_{k}^{l} + 1)^{n} \end{aligned}

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

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

k \times a n s 2 = \sum_{l = 0}^{k - 1} (p ω_{k}^{l} + 1)^{n} \sum_{j = 0}^{k - 1} j ω_{k}^{- j l}

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

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

\begin{aligned} S (x) = & \sum_{j = 0}^{k - 1} j ω_{k}^{j x} \\ ω_{k}^{x} S (x) = & \sum_{j = 1}^{k} (j - 1) ω_{k}^{j x} \\ (ω_{k}^{x} - 1) S (x) = & (k - 1) ω_{k}^{j k} - \sum_{j = 1}^{k - 1} ω_{k}^{j x} \end{aligned}

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

那么我们可以得到 $S (x) = \frac{k}{ω_{k}^{x} - 1}$ ，唯一注意 $n | x$ 的时候 $S (x) = \frac{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;
}