简述
在计数题解题中遇到要求 \(x \bmod m = 0\) 的条件时，如果式子的计算复杂度比较高，但是式子里含有二项式系数的时候，就可以考虑用单位根反演。
原理
长这个样子：
\[ \frac{1}{m} \sum_{i = 0}^{m – 1} (\omega_m^k)^i = [m | k] \]
证明不难，留作习题。
例题 A：LibreOJ #6358. 前夕
这道题题意：给 \(n\) 个元素，然后选出若干个这些元素的集合，使得这些集合的交集是 \(4\) 的倍数。
可以考虑先试试暴力容斥。枚举交集大小，然后乘上容斥系数 \(f(i)\)：
\[ \sum_{k = 0}^n f(i) {n \choose k} (2^{2^{n – k}} – 1) \]
考虑计算这个容斥系数，然后就可以算出结果了。这个容斥系数需要对前 \(n\) 个满足：
\[ \sum_{i = 0}^n {n \choose i} f(i) = [4 | n] \]
二项式反演一下：
\[ f(n) = \sum_{i = 0}^n (-1)^{n – i} {n \choose i} [4 | i] \]
正常计算要 \(\Theta(n^2)\) 的时间，肯定超时。设 \(m = 4\)，我们尝试套单位根反演进去：
\[ \begin{aligned} f(k) &= \sum_{i = 0}^k (-1)^{k – i} {k \choose i} [m | i] \\ &= \sum_{i = 0}^k (-1)^{k – i} {k \choose i} \frac{1}{m} \sum_{j = 0}^{m – 1} (\omega_m^i)^j \\ &= \frac{1}{m} \sum_{i = 0}^k \sum_{j = 0}^{m – 1} (-1)^{k – i} {k \choose i} (\omega_m^i)^j \\ &= \frac{1}{m} \sum_{i = 0}^{m – 1} \sum_{j = 0}^{k} {k \choose j} (\omega_m^i)^j (-1)^{k – j} \end{aligned} \]
然后就是喜闻乐见的二项式定理合并：
\[ \begin{aligned} f(k) &= \frac{1}{m} \sum_{i = 0}^{m – 1} \sum_{j = 0}^{k} {k \choose j} (\omega_m^i)^j (-1)^{k – j} \\ &= \frac{1}{m} \sum_{i = 0}^{m – 1} (\omega_m^i -1)^k \end{aligned} \]
然后直接算就完事了。最后，在这里简单说明一下单位根一般的写法：找到模数 \(P\) 的原根 \(g\) 之后，那么 \(\omega_m = \frac{g^{P – 1}}{m}\)。
// LOJ6358.cpp
#include <bits/stdc++.h>
using namespace std;
const int MAX_N = 1e7 + 200, m = 4, mod = 998244353, G = 3;
int f[MAX_N], n, fac[MAX_N], fac_inv[MAX_N];
int quick_pow(int bas, int tim, int cmod)
{
int ret = 1;
while (tim)
{
if (tim & 1)
ret = 1LL * ret * bas % cmod;
bas = 1LL * bas * bas % cmod;
tim >>= 1;
}
return ret;
}
const int wn = quick_pow(G, (mod - 1) / 4, mod);
int wns[4], wn_org[4];
int binomial(int n_, int k_) { return 1LL * fac[n_] * fac_inv[k_] % mod * fac_inv[n_ - k_] % mod; }
int main()
{
scanf("%d", &n);
for (int i = fac[0] = fac_inv[0] = 1; i <= n; i++)
fac[i] = 1LL * fac[i - 1] * i % mod;
fac_inv[n] = quick_pow(fac[n], mod - 2, mod);
for (int i = n - 1; i >= 1; i--)
fac_inv[i] = 1LL * fac_inv[i + 1] * (i + 1) % mod;
for (int i = 0; i < 4; i++)
wn_org[i] = quick_pow(wn, i, mod) - 1, wns[i] = 1;
int m_inv = quick_pow(m, mod - 2, mod);
for (int k = 0; k <= n; k++)
{
for (int i = 0; i < m; i++)
f[k] = (1LL * f[k] + wns[i]) % mod;
f[k] = 1LL * f[k] * m_inv % mod;
for (int i = 0; i < m; i++)
wns[i] = 1LL * wns[i] * wn_org[i] % mod;
}
int ans = 0;
for (int k = n, pow_2 = 2; k >= 0; k--, pow_2 = 1LL * pow_2 * pow_2 % mod)
{
// subset is k;
int tmp = 1LL * f[k] * binomial(n, k) % mod * ((pow_2 + mod - 1) % mod) % mod;
ans = (0LL + ans + tmp) % mod;
}
printf("%d\n", ans + 1);
return 0;
}
例题 B：BZOJ 3328 – PYXFIB
其实可以考虑放在 Fibonacci 的矩阵上做：矩阵满足单位根反演的规则。那么我们可以把式子写成：
\[ \begin{aligned} & \ \ \ \ \ \sum_{i = 0}^n {n \choose i} F_i [k | i] \\ &= \sum_{i = 0}^n {n \choose i} F_i \cdot \frac{1}{k} \sum_{j = 0}^{k – 1} (\omega_k^i)^j \\ &= \frac{1}{k} \sum_{i = 0}^n \sum_{j = 0}^{k – 1} {n \choose i} F_i (\omega_k^i)^j \\ &= \frac{1}{k} \sum_{i = 0}^{k – 1} \sum_{j = 0}^{n} {n \choose j} F_j \omega_k^{ij} \end{aligned} \]
然后我们可以把 Fibonacci 的部分换成转移矩阵 \(A\)：
\[ \begin{aligned} & \ \ \ \ \ \frac{1}{k} \sum_{i = 0}^{k – 1} \sum_{j = 0}^{n} {n \choose j} A^j \omega_k^{ij} \\ &= \frac{1}{k} \sum_{i = 0}^{k – 1} \sum_{j = 0}^{n} {n \choose j}I^{n – j} (A \omega_k^i)^j \\ &= \frac{1}{k} \sum_{i = 0}^{k – 1} (A \omega_k^i + I)^n \end{aligned} \]
然后就完事了。
// BZ3328.cpp
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int MAX_N = 1e5 + 200;
int T, m, mod, ptot, primes[MAX_N], g, wns[MAX_N];
ll n;
struct matrix
{
int mat[2][2];
void clear() { memset(mat, 0, sizeof(mat)); }
int *operator[](const int &rhs) { return mat[rhs]; }
matrix operator+(const matrix &rhs)
{
matrix ret;
for (int i = 0; i < 2; i++)
for (int j = 0; j < 2; j++)
ret[i][j] = (1LL * mat[i][j] + rhs.mat[i][j]) % mod;
return ret;
}
matrix operator*(const matrix &rhs)
{
matrix ret;
ret.clear();
for (int i = 0; i < 2; i++)
for (int j = 0; j < 2; j++)
for (int k = 0; k < 2; k++)
ret[i][j] = (1LL * ret[i][j] + 1LL * mat[i][k] * rhs.mat[k][j] % mod) % mod;
return ret;
}
matrix operator*(const int &rhs)
{
matrix ret;
for (int i = 0; i < 2; i++)
for (int j = 0; j < 2; j++)
ret[i][j] = 1LL * mat[i][j] * rhs % mod;
return ret;
}
matrix operator^(const ll &rhs);
} eps, fib;
matrix matrix::operator^(const ll &rhs)
{
ll tim = rhs;
matrix ret = eps, bas = *this;
while (tim)
{
if (tim & 1LL)
ret = ret * bas;
bas = bas * bas;
tim >>= 1;
}
return ret;
}
int quick_pow(int bas, int tim, int cmod)
{
int ret = 1;
while (tim)
{
if (tim & 1)
ret = 1LL * ret * bas % cmod;
bas = 1LL * bas * bas % cmod;
tim >>= 1;
}
return ret;
}
void find_root()
{
ptot = 0;
int x = mod - 1;
for (int i = 2; 1LL * i * i <= x; i++)
if (x % i == 0)
{
primes[++ptot] = i;
while (x % i == 0)
x /= i;
}
if (x > 1)
primes[++ptot] = x;
for (int i = 2; i <= mod - 1; i++)
{
bool flag = true;
for (int k = 1; flag && k <= ptot; k++)
if (quick_pow(i, (mod - 1) / primes[k], mod) == 1)
flag = false;
if (flag)
{
g = i;
break;
}
}
int wn = quick_pow(g, (mod - 1) / m, mod);
wns[0] = 1;
for (int i = 1; i < m; i++)
wns[i] = 1LL * wns[i - 1] * wn % mod;
}
int main()
{
for (int i = 0; i < 2; i++)
eps.mat[i][i] = 1;
fib[0][0] = fib[0][1] = fib[1][0] = 1;
scanf("%d", &T);
while (T--)
{
scanf("%lld%d%d", &n, &m, &mod), find_root();
matrix ans;
ans.clear();
for (int i = 0; i < m; i++)
ans = ans + ((fib * wns[i] + eps) ^ n);
printf("%lld\n", 1LL * ans[0][0] * quick_pow(m, mod - 2, mod) % mod);
}
return 0;
}