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

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 整除、下标同余，直接套单位根反演。