单位根反演
WorldMachine
·
·
算法·理论
首先我们有式子：
[n\mid m]=\dfrac1n\sum_{i=0}^{n-1}\omega_n^{im}
证明：设 f(x)=\sum\limits_{i=0}^{n-1}x^i，那么原式即为 f(\omega_n^m)。当 n\mid m 时，求和的每一项都是 1，显然成立；否则有 f(x)=\dfrac{x^n-1}{x-1} 成立，由于 \omega_n^{nm}=1，故上式为 0。
这个式子可以扩展到 m\bmod n=k 的情况：
[m\bmod n=k]=[n\mid(m-k)]=\dfrac1n\sum_{i=0}^{n-1}\omega_n^{-ik}\omega_n^{im}
现在有多项式 f(x)=\sum\limits_{i=0}^{n}a_ix^i，要求其所有下标模 m 等于 k 的系数之和，有：
\sum\limits_{i=0}^na_i[i\bmod m=k]=\dfrac1m\sum_{i=0}^na_i\sum_{j=0}^{m-1}\omega_m^{-jk}\omega_m^{ij}
交换求和符号：
\dfrac1m\sum_{i=0}^{m-1}\omega_m^{-ik}\sum_{j=0}^na_i\omega_m^{ij}=\dfrac1m\sum_{i=0}^{m-1}\omega_m^{-ik}f(\omega_m^i)
这个式子的优势在于，如果多项式 f 的次数非常大，但它的点值可以在 \mathcal O(T) 的时间复杂度内计算，就可以用上式在 \mathcal O(mT) 的时间复杂度内求出 f 的所有下标模 m 为 k 的系数之和。
举个例子，现在要求：
\sum_{i=0}^n[i\bmod m=k]\binom ni\bmod p
$$
\dfrac1m\sum_{i=0}^{m-1}\omega^{-ik}(1+\omega_m^i)^n
$$
如果 $p\bmod m=1$，那么找到 $p$ 的原根 $g$，$g^{\frac{p-1}{m}}$ 即可作为 $\omega_m$ 运算，复杂度为 $\mathcal O(m\log n)$。
### [LOJ6485 LJJ 学二项式定理](https://loj.ac/p/6485)
$\dbinom{n}{i}s^i$ 的 OGF 即为 $(1+sx)^n$，直接用上面的推论计算即可。
```cpp
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int p = 998244353, g = 3;
int T, s, a[4], w[4]; ll n;
int qpow(int a, int b) {
int c = 1;
while (b) { if (b & 1) c = (ll)c * a % p; a = (ll)a * a % p, b >>= 1; }
return c;
}
int main() {
w[0] = 1;
const int W = qpow(g, (p - 1) / 4), iv4 = qpow(4, p - 2);
for (int i = 1; i < 4; i++) w[i] = (ll)w[i - 1] * W % p;
scanf("%d", &T);
while (T--) {
scanf("%lld%d%d%d%d%d", &n, &s, &a[0], &a[1], &a[2], &a[3]);
n %= p - 1;
int ans = 0;
for (int i = 0; i < 4; i++) {
int sum = 0;
for (int j = 0; j < 4; j++) {
sum = (sum + (ll)qpow(((ll)s * w[j] + 1) % p, n) * qpow(w[j], ll(p - 2) * i % (p - 1))) % p;
}
ans = (ans + (ll)sum * a[i] % p * iv4) % p;
}
printf("%d\n", ans);
}
}
```
### [P10664 BZOJ3328 PYXFIB](https://www.luogu.com.cn/problem/P10664)
所求即为：
$$
\sum_{i=0}^n[k\mid i]\binom ni\text{fib}_i
$$
使用单位根反演：
$$
\dfrac1k\sum_{i=0}^{k-1}\sum_{j=0}^n\omega_k^{ij}\binom nj\text{fib}_j
$$
这个斐波那契数在这里很难受，我们可以把它变成它的转移矩阵 $F=\begin{bmatrix}1&1\\0&1\end{bmatrix}$，现在要求：
$$
\dfrac1k\sum_{i=0}^{k-1}\sum_{j=0}^n\omega_k^{ij}\binom njF^j
$$
第二个求和式是二项式定理的形式，注意矩阵的二项式定理只在可交换矩阵之间适用，当然 $I$ 和任意同型矩阵都是可交换的，那原式其实就是 $(I+\omega_k^jF)^n$，用矩阵快速幂优化即可。
这个题难受之处在于要自己找原根，用随机找原根再 check 的方法即可。
```cpp
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
int T, k, p, ph, t, fc[15]; ll n;
mt19937 rnd(time(0));
int qpow(int a, int b) {
int c = 1;
while (b) { if (b & 1) c = (ll)c * a % p; a = (ll)a * a % p, b >>= 1; }
return c;
}
bool chk(int g) {
for (int i = 1; i <= t; i++) if (qpow(g, (p - 1) / fc[i]) == 1) return 0;
return 1;
}
struct mat { int a[2][2]; mat() { memset(a, 0, sizeof(a)); } };
mat operator*(const mat &a, const mat &b) {
mat c;
c.a[0][0] = ((ll)a.a[0][0] * b.a[0][0] + (ll)a.a[0][1] * b.a[1][0]) % p;
c.a[0][1] = ((ll)a.a[0][0] * b.a[0][1] + (ll)a.a[0][1] * b.a[1][1]) % p;
c.a[1][0] = ((ll)a.a[1][0] * b.a[0][0] + (ll)a.a[1][1] * b.a[1][0]) % p;
c.a[1][1] = ((ll)a.a[1][0] * b.a[0][1] + (ll)a.a[1][1] * b.a[1][1]) % p;
return c;
}
mat operator^(mat a, ll b) {
mat c; c.a[0][0] = c.a[1][1] = 1;
while (b) { if (b & 1) c = c * a; a = a * a, b >>= 1; }
return c;
}
int main() {
scanf("%d", &T);
while (T--) {
scanf("%lld%d%d", &n, &k, &p), ph = p - 1, t = 0;
for (int i = 2; i * i <= ph; i++) {
if (ph % i) continue;
fc[++t] = i; while (!(ph % i)) ph /= i;
}
if (ph != 1) fc[++t] = ph;
int g = 0;
while (1) { if (g && chk(g)) break; g = rnd() % (p - 1) + 1; }
int w = qpow(g, (p - 1) / k), ans = 0;
for (int i = 0; i < k; i++) {
mat a; a.a[0][0] = a.a[0][1] = a.a[1][0] = qpow(w, i);
a.a[0][0]++, a.a[1][1]++, a = a ^ n, ans = (ans + a.a[0][0]) % p;
}
cout << (ll)ans * qpow(k, p - 2) % p << '\n';
}
}
```
### [P5591 小猪佩奇学数学](https://www.luogu.com.cn/problem/P5591)
想办法把这个向下取整变成可以单位根反演的形式。注意有 $\left\lfloor\dfrac ik\right\rfloor=\sum\limits_{j=0}^i[k\mid j]-1$，那么可以得到：
$$
\left\lfloor\dfrac ik\right\rfloor=\sum_{j=0}^i\dfrac 1k\sum_{l=0}^{k-1}\omega_k^{jl}-1
$$
所求即为：
$$
\sum_{i=0}^n\binom nip^i\left(\sum_{j=0}^i\dfrac 1k\sum_{l=0}^{k-1}\omega_k^{jl}-1\right)=\dfrac1k\sum_{i=0}^n\binom nip^i\sum_{j=0}^i\sum_{l=0}^{k-1}\omega_{k}^{jl}-(1+p)^n
$$
$(1+p)^n$ 直接算就是了，考虑前面那一坨。交换求和符号：
$$
\sum_{l=0}^{k-1}\sum_{i=0}^n\binom nip^i\sum_{j=0}^n\omega_k^{jl}=\sum_{l=0}^{k-1}\sum_{i=0}^n\binom nip^i\cdot\dfrac{\omega_k^{(i+1)l}-1}{\omega_k^l-1}
$$
上面的式子其实有个问题，要注意等比数列求和公式的使用是有条件的，要求 $k\nmid l$，在求和范围的限制下也就是 $l\not=0$，因此 $l=0$ 的时候要特判，它等于：
$$
\sum_{i=0}^n\binom nip^i(i+1)=\sum_{i=0}^n\binom niip^i+\sum_{i=0}^n\binom nip^i=np(p+1)^{n-1}+(p+1)^n
$$
回到上面的式子，继续化简：
$$
\sum_{l=1}^{k-1}\dfrac{1}{\omega_k^l-1}\sum_{i=0}^n\binom nip^i(\omega_k^{(i+1)l}-1)
$$
把右边那一坨拆开：
$$
\omega_k^l\sum_{i=0}^n\binom nip^i\omega_k^{il}-\sum_{i=0}^n\binom nip^i=\omega_k^l(1+p\omega_k^l)^n-(1+p)^n
$$
这样就可以用快速幂计算了。
```cpp
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int mod = 998244353, g = 3;
int n, p, k, w;
int qpow(int a, int b) {
int c = 1;
while (b) { if (b & 1) c = (ll)c * a % mod; a = (ll)a * a % mod, b >>= 1; }
return c;
}
int main() {
cin >> n >> p >> k, w = qpow(g, (mod - 1) / k);
int ans = ((ll)n * p % mod * qpow(1 + p, n - 1) + qpow(1 + p, n)) % mod;
for (int i = 1, pd = 1; i < k; i++) {
pd = (ll)pd * w % mod;
ans = (ans + ((ll)pd * qpow(1 + (ll)pd * p % mod, n) - qpow(1 + p, n) + mod) % mod * qpow(pd - 1, mod - 2)) % mod;
}
ans = ((ll)ans * qpow(k, mod - 2) - qpow(1 + p, n) + mod) % mod;
cout << ans;
}
```