Number Theory
Data Structures & Algorithms

NTT

Polynomial convolution modulo 998244353 using roots of unity and iterative butterfly layers.

Category Number Theory
Level advanced
Source TeX + C++
convolutionpolynomialsmodular arithmetic

NTT

The main note is rendered from the TeX source. Code lives in a separate C++ file so the write-up stays readable.

Overview

The number theoretic transform is the modular counterpart of FFT. In contest terms, it is the standard exact method for multiplying polynomials in \(O(n \log n)\) time when the coefficients live modulo a suitable prime.

Whenever a counting problem secretly asks for convolution, NTT is often the cleanest fast solution if the modulus fits.

When to Use It

Use it when:

  • you need polynomial convolution faster than \(O(n^2)\),

  • coefficients are naturally taken modulo a prime such as \(998244353\),

  • exact modular arithmetic is preferable to floating-point FFT,

  • the problem can be reframed as counting pair sums, subset combinations, or polynomial products.

  • For tiny input sizes, NTT is overkill. For incompatible moduli, a different transform strategy or CRT layer may be needed.

Core Idea

Convolution becomes pointwise multiplication after transforming to the right basis. FFT uses complex roots of unity. NTT does the same thing inside modular arithmetic using roots of unity modulo a prime.

For the common modulus \(998244353\), \[ 998244353 - 1 = 119 \cdot 2^{23}, \] so there are enough \(2^k\)-th roots of unity for transform lengths up to \(2^{23}\). That is exactly why this modulus appears so often in contests.

Key Insight

The transform is fast because the evaluation points are highly structured. Instead of evaluating a polynomial at every point independently, the butterfly layers reuse work between even and odd parts.

The other key fact to remember is that the inverse transform is not ``free''. After reversing the butterfly process, every coefficient must be multiplied by \(n^{-1} \bmod \texttt{MOD}\).

Operations / Main Technique

The standard convolution pipeline is:

  • choose a power-of-two length \(n\) large enough for the result,

  • pad both arrays to length \(n\),

  • run the forward transform on both,

  • multiply pointwise,

  • run the inverse transform,

  • normalize by \(n^{-1}\).

Worked viewpoint.

If \(A(x)\) and \(B(x)\) are transformed into their values at the same roots of unity, then \[ (A \cdot B)(\omega_i) = A(\omega_i) B(\omega_i). \] That is the whole reason pointwise multiplication in transform space corresponds to convolution in coefficient space.

Correctness Intuition

The chosen roots of unity make the transform matrix invertible, just like the complex DFT. So the forward transform is an evaluation map at special points, and the inverse transform reconstructs the unique polynomial coefficients from those values.

The butterfly layers are only a structured implementation of that algebra. Pointwise multiplication after the forward transform gives the values of the product polynomial at those same points, and the inverse transform recovers the coefficients.

Complexity Analysis

  • forward transform: \(O(n \log n)\),

  • inverse transform: \(O(n \log n)\),

  • convolution: \(O(n \log n)\) total after padding,

  • memory: \(O(n)\).

  • Here \(n\) is the padded transform length, not the original polynomial degree.

Implementation

The reference code uses:

  • modulus \(998244353\),

  • primitive root \(3\),

  • iterative bit reversal,

  • a convolution(a, b) wrapper.

  • That version is the contest-ready baseline I want to understand before hiding everything behind a library call.

Common Pitfalls

  • Using a modulus that does not support the required transform length.

  • Forgetting inverse normalization by \(n^{-1}\).

  • Padding to a length smaller than \(|a| + |b| - 1\).

  • Mixing negative values into modular arithmetic without re-normalizing.

  • Using NTT when a quadratic convolution would have been simpler and fast enough.

Variants / Extensions

  • CRT-based convolution for arbitrary moduli,

  • AtCoder Library convolution wrappers,

  • formal power series routines built on repeated NTT,

  • subset of problems where FFT with doubles is acceptable instead.

  • For most CP uses, direct polynomial multiplication is only the entry point. Once NTT is available, many generating function and counting tricks become feasible.

Practice Problems

  • Multiply two polynomials modulo \(998244353\).

  • Count pair sums by multiplying frequency polynomials.

  • DP transitions that become convolution after regrouping states.

References

Code

Contest-ready reference implementation for the idea explained above.

C++ competitive_programming/dsa/number-theory/ntt/code.cpp

Kept as a standalone source file so the implementation can be copied without TeX markup around it.

Raw file
namespace ntt998244353 {
const int MOD = 998244353;
const int PRIMITIVE_ROOT = 3;

int mod_pow(int base, long long exponent) {
    int result = 1;
    while (exponent > 0) {
        if (exponent & 1) result = (long long)result * base % MOD;
        base = (long long)base * base % MOD;
        exponent >>= 1;
    }
    return result;
}

void ntt(vector<int>& a, bool invert) {
    int n = (int)a.size();

    for (int i = 1, j = 0; i < n; ++i) {
        int bit = n >> 1;
        while (j & bit) {
            j ^= bit;
            bit >>= 1;
        }
        j ^= bit;
        if (i < j) swap(a[i], a[j]);
    }

    for (int len = 2; len <= n; len <<= 1) {
        int wlen = mod_pow(PRIMITIVE_ROOT, (MOD - 1) / len);
        if (invert) wlen = mod_pow(wlen, MOD - 2);

        for (int start = 0; start < n; start += len) {
            long long w = 1;
            for (int i = 0; i < len / 2; ++i) {
                int u = a[start + i];
                int v = (int)(w * a[start + i + len / 2] % MOD);
                a[start + i] = u + v < MOD ? u + v : u + v - MOD;
                a[start + i + len / 2] = u - v >= 0 ? u - v : u - v + MOD;
                w = w * wlen % MOD;
            }
        }
    }

    if (invert) {
        int inv_n = mod_pow(n, MOD - 2);
        for (int& x : a) x = (long long)x * inv_n % MOD;
    }
}

vector<int> convolution(vector<int> a, vector<int> b) {
    if (a.empty() || b.empty()) return {};

    int needed = (int)a.size() + (int)b.size() - 1;
    int n = 1;
    while (n < needed) n <<= 1;

    a.resize(n);
    b.resize(n);
    ntt(a, false);
    ntt(b, false);
    for (int i = 0; i < n; ++i) {
        a[i] = (long long)a[i] * b[i] % MOD;
    }
    ntt(a, true);
    a.resize(needed);
    return a;
}
}  // namespace ntt998244353

Source Files and Assets

Raw files are still available here when you want the original TeX, C++, or statement assets.

Show raw files