C++ 单位根反演（Roots of Unity Filter）全解析

前言

单位根反演是组合数学、数论、生成函数、多项式算法中最核心的等价变换公式，专门解决整除性判定、模 k 计数、下标取模筛选、周期性求和问题。它能把「对下标 i ≡ r (mod k) 求和」这种带取模限制的求和，转化为无限制普通求和，从而大幅简化计算。

在算法竞赛（ACM/NOI/CSP）、密码学、信号处理、C++ 高性能数学库、NTT 相关算法里，单位根反演是必备工具。它的本质是：用复数/模意义下单位根，构造“筛选器”，精准提取满足同余条件的项。

本文从数学定义→公式推导→几何意义→经典题型→C++ 完整实现全覆盖，理论严谨、代码可直接运行。

---

目录

1. 单位根基础（复数域 + 模素数域）
2. 单位根反演公式核心推导
3. 公式直观理解：筛选器原理
4. 最常用 3 大模板公式
5. 典型应用场景（取模求和、整除计数、循环卷积）
6. C++ 完整实现（复数版 + 模素数 NTT 版）
7. 例题解析与代码运行演示
8. 复杂度与工程优化

---

一、前置知识：n 次单位根

1.1 复数域单位根

n 次单位根：满足 ( z^n = 1 ) 的复数。
主 n 次单位根：
[
\omega_n = e^{\frac{2\pi i}{n}} = \cos\frac{2\pi}{n} + i\sin\frac{2\pi}{n}
]

核心性质：

1. ( \omega_n^n = 1 )
2. ( \omega_n^k ) 两两不同
3. 正交性（反演基础）：
   [
   \frac{1}{n}\sum_{j=0}^{n-1} \omega_n^{j\cdot k} = \begin{cases} 1 & n\mid k \ 0 & else \end{cases}
   ]

1.2 模素数域单位根（算法中最常用）

算法里不用复数，用模素数域，通过原根构造单位根：
设模数 ( P ) 是质数，( g ) 是原根，则：
[
\omega_n = g^{\frac{P-1}{n}}
]
要求：n | (P-1)。

常用模数：

• 998244353（原根3）
• 1004535809（原根3）

---

二、单位根反演核心公式（最重要）

2.1 基础公式（正交性）

[
\boxed{\frac{1}{k} \sum_{j=0}^{k-1} \omega_k^{j\cdot t} = [k \mid t]}
]

含义：

• 若 t 是 k 的倍数 → 结果 = 1
• 否则 → 结果 = 0

这就是整除判定器。

2.2 反演公式（求和筛选）

我们要计算：
[
S = \sum_{n} [k \mid n] \cdot a_n
]

代入单位根反演：
[
S = \frac{1}{k} \sum_{j=0}^{k-1} \sum_{n} a_n \cdot \omega_k^{jn}
]

交换求和顺序：
[
S = \frac{1}{k} \sum_{j=0}^{k-1} F(\omega_k^j)
]
其中 ( F(x) = \sum a_n x^n ) 是生成函数。

一句话总结：单位根反演 = 把取模限制，变成对单位根的求值。

2.3 更一般：同余 n ≡ r (mod k)

[
\sum_{n \equiv r \pmod k} a_n
= \frac{1}{k}\sum_{j=0}^{k-1} \omega_k^{-jr} F(\omega_k^j)
]

这是算法竞赛最常用模板。

---

三、直观理解：单位根 = 筛选器

你可以把单位根反演看作：

• 构造 k 个“频率”
• 只有下标是 k 倍数的项会共振保留
• 其余项全部抵消为 0

它就是数学里的带通滤波器。

---

四、单位根反演三大标准模板（背会直接用）

模板 1：求下标被 k 整除的和

[
\sum_{k\mid n} a_n
= \frac1k \sum_{j=0}^{k-1} F(\omega_k^j)
]

模板 2：求下标 ≡ r (mod k) 的和

[
\sum_{n\equiv r\pmod k} a_n
= \frac1k \sum_{j=0}^{k-1} \omega_k^{-jr} F(\omega_k^j)
]

模板 3：计数 f(n) 中满足 n mod k = r 的数量

[
ans = \frac1k \sum_{j=0}^{k-1} \omega_k^{-jr} \sum_n f(n)\omega_k^{jn}
]

---

五、典型应用场景

1. 求 n! 中指数是 k 的倍数的项和
2. 斐波那契数列中满足下标 ≡0/1/2 mod k 的和
3. 组合数取模筛选
4. 循环卷积、周期序列求和
5. 生成函数限定下标提取
6. NTT 加速的取模类 DP 优化

凡是题目出现：
“mod k”、“被 k 整除”、“下标同余”、“周期为 k”
99% 可以用单位根反演。

---

六、C++ 完整实现

版本 1：复数版（教学、直观）

#include <iostream>
#include <vector>
#include <complex>
#include <cmath>
using namespace std;

typedef complex<double> cd;
const double PI = acos(-1);

// 主 n 次单位根
cd omega(int n, int k) {
    return exp(cd(0, 2 * PI * k / n));
}

// 单位根反演：求 a_n 中 n ≡ r (mod k) 的和
template<typename Func>
double roots_of_unity_filter(Func&& F, int k, int r) {
    cd res = 0;
    for (int j = 0; j < k; j++) {
        cd w = omega(k, j);
        cd coeff = exp(cd(0, -2 * PI * j * r / k)); // w^(-jr)
        res += coeff * F(w);
    }
    return res.real() / k;
}

// 示例：F(x) = 1 + x + x^2 + x^3 + x^4 + x^5
double test_F(cd x) {
    return 1.0 + x + x*x + x*x*x + x*x*x*x + x*x*x*x*x;
}

int main() {
    int k = 3, r = 0;
    double ans = roots_of_unity_filter(test_F, k, r);
    // 下标 0,3 → 1+1=2
    cout << "n ≡ 0 mod 3 的和 = " << ans << endl;
    return 0;
}

版本 2：模素数版（算法竞赛正式版）

真正比赛/工程用这个

#include <iostream>
#include <vector>
using namespace std;

const int MOD = 998244353;
const int G = 3;

long long qpow(long long a, long long b) {
    long long res = 1;
    while (b) {
        if (b & 1) res = res * a % MOD;
        a = a * a % MOD;
        b >>= 1;
    }
    return res;
}

// 求 k 次单位根
long long get_omega(int k) {
    return qpow(G, (MOD - 1) / k);
}

// 单位根反演求 sum [n ≡ r mod k] a_n
long long solve(const vector<long long>& a, int k, int r) {
    long long w = get_omega(k);
    long long inv_k = qpow(k, MOD - 2);
    long long ans = 0;

    for (int j = 0; j < k; j++) {
        long long wj = qpow(w, j);
        long long c = qpow(wj, (k - r) % k); // w^(-jr)

        long long F = 0;
        long long p = 1;
        for (auto v : a) {
            F = (F + v * p) % MOD;
            p = p * wj % MOD;
        }

        ans = (ans + c * F) % MOD;
    }
    ans = ans * inv_k % MOD;
    return ans;
}

int main() {
    vector<long long> a = {1,1,1,1,1,1}; // 1+x+x^2+x^3+x^4+x^5
    int k = 3, r = 0;
    cout << solve(a, k, r) << endl; // 输出 2
    return 0;
}

---

七、公式推导（最关键的一步）

我们要证明：
[
\frac1k\sum_{j=0}^{{k-1}\omega_k}{jn} = [k\mid n]
]

证明：
设 ( n = qk + t )
若 ( t=0 )：
[
\frac1k\sum_{j=0}^{k-1} 1 = 1
]

若 ( t≠0 )：
[
\frac1k \cdot \frac{\omega_k^{{kn}-1}{\omega_k}n-1} = \frac1k \cdot 0 = 0
]

证毕。

---

八、复杂度分析

• 暴力筛选：( O(n) )
• 单位根反演：( O(k \cdot M) )
  • M 是生成函数求值代价
• 若配合 NTT：可做到 ( O(n\log n) )

k 通常很小（如 2,3,4,5），速度极快。

---

九、适用条件

1. 必须存在 k 次单位根
2. 模意义下要求 ( k \mid (P-1) )
3. 生成函数可快速求值
4. 问题包含同余、整除、周期

---

十、总结（最核心 3 句话）

1. 单位根反演 = 用单位根构造整除筛选器
2. 核心公式：
   [
   \frac1k\sum_{j=0}^{{k-1}\omega_k}{jn} = [k\mid n]
   ]
3. 凡是题目出现 mod k、被 k 整除、下标同余，直接套单位根反演。