NTT实战：从序列统计问题看模意义下多项式快速幂与循环卷积

1. 项目概述从一道经典题看NTT的实战威力最近在整理算法竞赛的笔记翻到了这道“jzoj4051-序列统计”。这道题在圈内挺有名的算是多项式算法特别是NTT快速数论变换的一个经典应用场景。很多朋友第一次接触NTT可能都是从FFT快速傅里叶变换学起的知道它能加速多项式乘法但一到具体题目尤其是模数下的乘法就有点发怵。这道题恰好就是一个绝佳的“练兵场”它把序列计数、生成函数和NTT紧密地结合在了一起让你能清晰地看到一个看似复杂的组合计数问题是如何被优雅地转化为多项式运算并被高效解决的。简单来说题目给你一个元素范围在[0, M-1]的集合S问你有多少种方式能构造出一个长度为N的序列使得序列中每个元素都属于S并且序列中所有元素的乘积模M等于一个给定的值X。这里的N可以非常大比如10^9级别M是一个质数且通常不大比如10^3级别。暴力枚举序列|S|^N的复杂度想都别想。动态规划N这么大状态转移也吃不消。这时候多项式特别是利用NTT进行加速的多项式幂就成了破题的关键。我之所以想详细聊聊这道题是因为它不仅仅是一个算法模板的应用。在解这道题的过程中你会深刻理解生成函数如何将组合意义“翻译”成代数形式理解原根在模意义下如何扮演单位根的角色从而使得FFT的思想得以在整数模域中复用最终理解NTT如何将理论上的O(n^2)乘法优化到O(n log n)。更重要的是你会掌握一套解决“带模数约束的序列计数”问题的通用思路。下面我就把自己琢磨这道题时的完整思考路径、实现细节和踩过的坑系统地分享出来。2. 核心思路拆解从组合问题到多项式乘法面对一个复杂问题第一步永远是化繁为简寻找其数学本质。这道题的核心约束有两个1) 序列长度固定为N2) 序列元素的乘积模M为定值X。我们的武器库里有生成函数。2.1 生成函数的引入与构造生成函数是连接离散组合与连续分析的桥梁。对于这类“从集合S中选取元素构成序列”的问题我们为每个元素a分配一个形式变量。但这里我们不能直接用x^a因为我们的最终目标是元素的乘积而非和。一个巧妙的技巧是利用指数。设M是一个质数根据数论知识存在原根g。这意味着[1, M-1]中的任意非零数都可以表示为g^k mod M的形式其中k在[0, M-2]之间。对于0这个特殊元素我们需要单独处理因为0没有离散对数。题目中集合S通常包含0这使得情况稍微复杂一些。我们先考虑S中非零元素的部分。对于任意一个非零元素s我们可以找到唯一的指数t使得s ≡ g^t (mod M)。这样序列中元素的乘积s1 * s2 * ... * sN mod M就等于g^(t1 t2 ... tN) mod M。原来对乘积取模的约束就转化为了对指数和取模(M-1)的约束基于这个转化我们可以为集合S构造一个生成函数A(x)A(x) Σ_{s∈S且s≠0} x^{ind_g(s)}其中ind_g(s)是s以g为底的离散对数即指数t。 同时我们记zero_flag表示0是否在集合S中。那么构造一个长度为N的序列其所有非零元素对应的指数和为T的方案数就等于生成函数[A(x)]^N中x^T项的系数。这里的多项式乘法对应着指数和的组合。2.2 目标系数的提取与零元素的处理我们的目标是序列乘积模M等于X。如果X 0那么序列中只要出现至少一个0即可。方案数 总序列数 - 全由非零元素构成的序列数。总序列数 |S|^N。全非零序列数这相当于从S中剔除0后得到的集合S其生成函数为A(x)。全非零序列的生成函数是[A(x)]^N。所有全非零序列的总数就是[A(x)]^N的所有系数之和。由于A(x)的每一项都是1表示选取该指数对应的元素一次[A(x)]^N的所有系数之和等于(Σ_{i} [x^i]A(x))^N |S|^N。这个结果很直观。 因此当X0时答案 |S|^N - |S|^N。这里涉及模意义下的快速幂非常简单。如果X ! 0那么序列中不能包含0。设X ≡ g^K (mod M)。我们需要序列所有元素均为非零的指数和T满足T ≡ K (mod M-1)。 我们需要的是[A(x)]^N这个多项式中所有指数模(M-1)余K的项的系数之和。 如何高效计算[A(x)]^N并求和呢直接计算多项式的N次幂哪怕用NTT如果N很大复杂度也是O(log N * len log len)但这里len是多少A(x)的最高次可能到M-2但N次幂后多项式的次数会达到N*(M-2)这是无法接受的。2.3 关键技巧循环卷积与DFT/NTT这里就需要用到第二个关键技巧循环卷积。我们关心的并不是指数和T的具体值而是T mod (M-1)的值。在多项式乘法中如果我们让乘法在模(x^{(M-1)} - 1)的意义下进行那么结果多项式里x^T项的系数就是所有原始指数和模(M-1)为T的方案的系数之和。这正是我们需要的而离散傅里叶变换DFT或其数论版本NTT其本质就是在计算循环卷积。当我们用NTT计算两个长度为LL是2的幂且大于等于循环节长度的向量a和b的循环卷积时过程是计算A NTT(a),B NTT(b)。计算C[i] A[i] * B[i] mod P。计算c INTT(C)。 得到的c就是a和b的循环卷积其中c[k] Σ_{ij ≡ k (mod L)} a[i]*b[j]。看这个形式和我们的需求完全一致我们只需要将生成函数A(x)表示为一个长度为L的向量a其中L是大于等于M-1的最小2的幂。a[i] 1如果存在s∈S使得ind_g(s) mod L i否则a[i]0。注意因为指数t在[0, M-2]之间而L M-1所以直接按t的值放在对应位置即可不会冲突。计算向量a的N次幂在循环卷积意义下。这可以通过计算NTT(a)得到A然后计算A[i]^N mod P得到C最后对C做逆NTT得到c。最终答案就是c[K]因为K在[0, M-2]范围内且L M-2所以c[K]就是我们需要的所有指数和模(M-1)余K的系数和。这个思路将问题完美转化一个复杂的组合计数问题 → 构造一个特定向量 → 利用NTT计算该向量在循环卷积意义下的N次幂 → 读取目标位置的系数。时间复杂度为O(L log L L log N)其中L是O(M)级别的完全可接受。3. 算法实现细节与NTT模板解析思路清晰了接下来就是实现。实现分为几个核心部分原根寻找、离散对数映射、NTT模板、多项式快速幂。这里我结合自己的代码分享一下需要注意的细节。3.1 原根的寻找与离散对数表生成题目保证M是质数。对于一个质数M其原根g满足g^1, g^2, ..., g^{M-1} mod M恰好是1到M-1的一个排列。 寻找原根有一个简单的方法从小到大枚举g从2开始检查对于M-1的所有素因子p_i是否都有g^{(M-1)/p_i} ! 1 (mod M)。如果都成立则g是原根。// 寻找质数 mod 的原根 int find_primitive_root(int mod) { int phi mod - 1; // 欧拉函数对于质数就是 M-1 vectorint factors; int temp phi; // 质因数分解 phi for (int i 2; i * i temp; i) { if (temp % i 0) { factors.push_back(i); while (temp % i 0) temp / i; } } if (temp 1) factors.push_back(temp); for (int g 2; g mod; g) { bool ok true; for (int factor : factors) { if (qpow(g, phi / factor, mod) 1) { // 快速幂 ok false; break; } } if (ok) return g; } return -1; // 理论上质数一定有原根 }找到原根g后我们需要预处理一个ind数组使得ind[g^i mod M] i。这样对于集合S中的非零元素s我们可以用O(1)时间得到其指数ind[s]。vectorint ind(M, -1); int cur 1; for (int i 0; i M-1; i) { ind[cur] i; cur (cur * g) % M; } // 现在 ind[s] 就是 s 的离散对数 (s ! 0)3.2 NTT模板的选择与参数设置NTT是在模素数P下进行的DFT要求P可以表示为c * 2^k 1的形式且2^k要足够大至少大于我们需要的变换长度L。常见的NTT模数有998244353 (119*2^231)1004535809 (479*2^211)等。本题通常取P 1004535809因为它的原根是3且2^21足够处理M在1000量级的问题L需要是2的幂且 M-1M-1最大约1000所以L取2048足够远小于2^21。NTT模板需要实现正变换NTT和逆变换INTT。核心是蝴蝶操作和迭代实现。这里给出一个常用的迭代NTT实现const int P 1004535809; // 模数 const int G 3; // 原根 const int Gi 334845270; // 原根的逆元即 pow(G, P-2, P) int qpow(int a, int b) { /* 快速幂 */ } void ntt(vectorint a, int inv) { int n a.size(); // 比特位反转 for (int i 0, j 0; i n; i) { if (i j) swap(a[i], a[j]); for (int k n 1; (j ^ k) k; k 1); } for (int len 2; len n; len 1) { int wn qpow(inv 1 ? G : Gi, (P - 1) / len); for (int i 0; i n; i len) { int w 1; for (int j 0; j len / 2; j) { int u a[i j]; int v 1LL * a[i j len/2] * w % P; a[i j] (u v) % P; a[i j len/2] (u - v P) % P; w 1LL * w * wn % P; } } } if (inv -1) { int inv_n qpow(n, P - 2); for (int x : a) x 1LL * x * inv_n % P; } }注意逆变换inv -1后一定要乘以n的逆元这才是完整的逆变换。这是新手最容易忘记的一步忘记后结果会整体放大n倍导致答案错误。3.3 多项式快速幂的实现我们需要计算向量a的N次幂循环卷积意义下。利用NTT这变得非常高效对a做一次NTT得到点值表示A。将A的每个点值进行快速幂A[i] qpow(A[i], N)。注意这里的幂次N可能很大需要使用快速幂算法。对A做一次逆NTT得到结果向量c。int L 1; while (L M-1) L 1; // 确定变换长度大于等于循环节长度(M-1)的最小2的幂 vectorint a(L, 0); // 根据集合S填充a如果 s!0则 a[ind[s]] 1 for (int s : S) { if (s ! 0) { a[ind[s]] 1; } } // 计算NTT(a) ntt(a, 1); // 点值快速幂 for (int i 0; i L; i) { a[i] qpow(a[i], N); } // 逆变换回来 ntt(a, -1); // 此时 a[K] 就是答案对于 X ! 0 的情况这里有一个极其重要的细节我们是在模P(NTT模数) 下进行运算但题目最终的答案可能需要输出对另一个模数比如10^97取模的结果。这意味着我们不能直接用a[K]作为答案。因为a[K]是模P下的结果而P和题目要求的模数可能不同。正确的做法是用任意模数多项式乘法MTT的思想或者直接使用双模数NTT如1004535809和998244353然后用中国剩余定理CRT合并最后再对题目要求模数取模。但本题有一个更巧妙的性质我们向量a的初始值只有0和1经过N次幂后结果a[K]是一个组合计数它一定是整数并且理论上不会超过题目模数很多。在M较小~1000N较大时a[K]可能很大但我们可以通过使用两个不同的NTT模数分别计算然后用CRT合并出一个精确的大整数最后再取模。在实际竞赛中如果时间紧迫有时会冒险假设结果在long long范围内直接用一个大质数如1004535809计算最后输出时再取题目模数。但这并不严谨。稳妥起见我推荐实现双模数NTT。4. 完整解题流程与代码框架将上述所有步骤串联起来我们得到完整的解题流程。以下是基于双模数NTT的代码框架和逻辑。4.1 主算法流程读入数据N, M, X, |S|以及集合S。预处理原根与离散对数找到模M的一个原根g。预处理数组indind[g^i mod M] i。处理特殊情况X 0计算total pow_mod(|S|, N, MOD)。计算non_zero_cnt pow_mod(|S| - zero_flag, N, MOD)。zero_flag是S中是否有0答案ans (total - non_zero_cnt MOD) % MOD。输出并返回。处理X ! 0计算K ind[X]。如果X不在离散对数表中说明无解但题目通常保证有解确定NTT长度L使其为大于等于M-1的最小2的幂。初始化两个向量a1, a2长度均为L用于两个不同的NTT模数P11004535809,P2998244353。遍历集合S对于每个非零元素s令t ind[s]设置a1[t] a2[t] 1。双模数NTT快速幂对a1在模P1下做NTT得到A1对A1每个点值做N次幂模P1逆NTT得到c1。对a2在模P2下做同样的操作得到c2。此时c1[K]和c2[K]分别是模P1和P2下的结果。中国剩余定理合并使用CRT合并同余方程x ≡ c1[K] (mod P1)x ≡ c2[K] (mod P2)解出的x是一个精确的整数在P1*P2范围内。最终取模将合并后的大整数x对题目要求的模数如MOD1e97取模得到最终答案。4.2 核心代码片段与解释这里给出双模数NTT快速幂的核心部分和CRT合并的示例// 假设已经实现了 mod1_ntt, mod2_ntt 以及对应的快速幂 qpow1, qpow2 const int MOD 1000000007; // 题目要求模数 const int P1 1004535809, G1 3; const int P2 998244353, G2 3; int main() { // ... 读入数据预处理原根和ind ... if (X 0) { // 计算 zero_flag, total, non_zero_cnt ... cout (total - non_zero_cnt MOD) % MOD endl; return 0; } int K ind[X]; int L 1; while (L M-1) L 1; vectorint a1(L, 0), a2(L, 0); for (int s : S) { if (s ! 0) { int t ind[s]; a1[t] 1; a2[t] 1; } } // 模 P1 下的变换 mod1_ntt(a1, 1); for (int i 0; i L; i) a1[i] qpow1(a1[i], N); mod1_ntt(a1, -1); long long res1 a1[K]; // 模 P2 下的变换 mod2_ntt(a2, 1); for (int i 0; i L; i) a2[i] qpow2(a2[i], N); mod2_ntt(a2, -1); long long res2 a2[K]; // 中国剩余定理合并 res1 (mod P1) 和 res2 (mod P2) // 解方程: x ≡ res1 (mod P1), x ≡ res2 (mod P2) // 使用公式: x res1 k * P1, 其中 k (res2 - res1) * inv(P1 mod P2) mod P2 long long inv qpow2(P1 % P2, P2-2); // P1 在模 P2 下的逆元 long long k ((res2 - res1) % P2 P2) % P2 * inv % P2; long long exact_x res1 k * P1; // 精确解 // 对题目模数取模 long long ans exact_x % MOD; cout ans endl; return 0; }实操心得双模数NTT的CRT合并部分一定要小心处理负数。(res2 - res1) % P2可能为负需要加P2再取模。另外计算逆元时要确保是在正确的模数下计算这里是在模P2下求P1的逆元。5. 常见问题与调试技巧这道题实现起来细节很多很容易出错。下面是我在多次实现和调试中总结的一些常见坑点和排查技巧。5.1 原根求解失败或离散对数表错误问题程序运行结果完全不对或者对于小数据都出错。排查首先验证你找到的g是否真的是原根。可以写一个简单的测试程序生成g^1, g^2, ..., g^{M-1} mod M检查它们是否互不相同且覆盖1到M-1。检查离散对数表ind的生成逻辑。确保cur从1开始循环M-1次每次cur (cur * g) % M。打印出ind数组检查是否每个1到M-1的数都有对应的非负索引且0的索引是-1或特殊值。技巧对于小的M比如M5, 7, 11可以手算原根和离散对数表与程序输出对比。5.2 NTT结果异常答案偏大或混乱问题X ! 0时答案明显错误或者换了数据后答案不对。排查逆变换忘记乘逆元这是最高频的错误。检查你的ntt函数在inv -1的分支里是否对每个系数都乘上了n的逆元。变换长度L不足L必须严格大于等于M-1且是2的幂。如果L小于M-1循环卷积会混叠导致结果错误。建议打印L的值确认。模数混淆确保NTT过程中所有的运算包括蝴蝶操作中的乘法、加法、快速幂都在正确的模数P1或P2下进行。为两个模数分别写两套函数或使用模板是清晰的做法。数组未清零a1和a2向量在初始化后确保所有位置为0再根据集合S赋值1。未清零的位置可能有随机值。调试技巧构造一个极小规模的测试用例。例如令M3原根g2S{1,2}N2X1。手算序列有[1,1], [2,2]两种乘积模3均为1。答案应为2。程序计算ind[1]0,ind[2]1。向量a [1, 1]长度L2。NTT(a) - 点值平方 - INTT。手动模拟这个过程与程序中间结果对比能快速定位是NTT过程出错还是前后逻辑出错。5.3 双模数CRT合并出错问题单模数NTT在小数据时结果正确但大数据或随机数据时错误。排查逆元计算错误CRT合并公式k (res2 - res1) * inv(P1) mod P2中inv(P1)是P1在模P2意义下的逆元不是模P1下的逆元。务必用qpow(P1 % P2, P2-2, P2)来计算。中间结果溢出res1 k * P1可能超过long long范围大约9e18。P1和P2都是1e9量级k最大可达P2-1因此k * P1最大约1e18加上res1仍在long long安全范围内。但为了保险可以使用__int128或者在合并过程中就取模。符号处理(res2 - res1) % P2在C中可能得到负数。必须写成((res2 - res1) % P2 P2) % P2。验证方法可以尝试用三个模数如再加一个469762049进行三模数NTT然后用CRT合并其中两个看结果是否与双模数合并的结果一致。或者对于小数据直接使用Python的大整数计算精确结果进行比对。5.4 时间复杂度与优化点复杂度分析主要耗时在于NTT。我们进行了两次长度为L的NTT正逆以及L次点值快速幂。每次NTT复杂度O(L log L)快速幂O(L log N)。总复杂度O(L log L L log N)。L是O(M)级别完全足够。优化如果M-1本身就是2的幂那么L M-1是最优的。点值快速幂时N可能很大但模数P1和P2是固定的可以预处理原根的幂但通常直接调用快速幂函数即可因为L不大。对于X 0的情况直接使用公式计算避免了NTT是最快的。5.5 关于零元素的再思考在X ! 0的情况下我们的算法完全排除了0这是正确的。在X 0的情况下我们使用了容斥原理。这里有一个边界情况如果集合S中不包含0即zero_flag false那么X 0的情况答案应为0因为无法构造出乘积为0的序列。我们的公式|S|^N - |S|^N中此时|S| |S|结果为0也是正确的。最后这道题将数论原根、离散对数、组合数学生成函数、多项式算法NTT、循环卷积完美结合是一道质量极高的综合题。理解并实现它对你掌握“用多项式处理模意义下计数问题”的套路有极大的提升。在实际比赛中如果遇到类似的“序列计数且约束与乘积/和模某数相关”的问题都可以尝试向生成函数和NTT方向思考。