[Competitive Programming] Number Theoretic Transform

A running collection of NTT notes and the problems that need it — one section per idea, added to over time.

Contents

1. What NTT is, and why not FFT — the same butterfly with roots of unity borrowed from a finite field instead of \(\mathbb{C}\).
2. Convolution is polynomial multiplication — the only thing NTT computes; everything else is a rewrite into this shape.
3. Correlation is convolution with one side reversed — including the cyclic case, which is where most problems actually land.
4. The template — iterative NTT, conv, and cyclic_correlate.
5. ABC471 G — Caeser Syllables — counting vowel runs under every rotation of the alphabet. (AtCoder Beginner Contest 471)

What NTT is, and why not FFT

FFT computes the Discrete Fourier Transform fast. The DFT needs one ingredient: a primitive \(n\)-th root of unity, something \(\omega\) with \(\omega^n = 1\) and no smaller positive power equal to 1. Over the complex numbers that’s \(\omega = e^{2\pi i/n}\), and everything works.

The NTT is the observation that \(\mathbb{C}\) was never load-bearing. The FFT derivation only ever uses that \(\omega\) is a primitive \(n\)-th root of unity in a commutative ring where \(n\) is invertible. Any field with such an \(\omega\) will do — in particular \(\mathbb{Z}/p\mathbb{Z}\) for a well-chosen prime \(p\). Same recursion, same butterflies, same \(O(n \log n)\); only the arithmetic changes from complex to int mod p.

The catch is that the root has to exist. \((\mathbb{Z}/p\mathbb{Z})^\times\) is cyclic of order \(p - 1\), so it contains an element of order \(n\) iff \(n \mid p - 1\). Since FFT-style algorithms want \(n\) to be a power of two, you want \(p - 1\) to be divisible by a large power of two. That is the entire reason for the famous constant:

\(2^{23} \mid p - 1\), so any transform length up to \(2^{23} = 8{,}388{,}608\) has a root. The primitive root of the field is \(g = 3\), so for a length-\(n\) transform you take

which has order exactly \(n\). A few other primes people keep in their notebook, all with primitive root 3:

Why competitive programmers reach for NTT over FFT. FFT with double is approximate. The result gets rounded back to an integer at the end, and that rounding is only safe while the true coefficients stay well under the ~53 bits of mantissa — coefficients of size \(C\) convolved at length \(n\) push intermediate magnitudes to roughly \(n C^2\), and somewhere past \(10^{14}\) or so you start losing the round. NTT has no such failure mode: every intermediate is an exact residue. The price is that you get the answer mod p, which matters in two cases:

• The problem already asks for the answer mod 998244353 — then there’s nothing to pay, which is exactly why so many problems name that modulus.
• You want an exact integer that might exceed \(p\) — then run three different NTT primes and CRT the results back.

The happy middle case, easy to miss: the answer is an exact integer but you can prove it’s smaller than \(p\). Then one prime is enough and no CRT is needed. The problem below is that case.

Convolution is polynomial multiplication

Everything NTT does is this one operation:

Read \(f\) and \(g\) as coefficient vectors of polynomials and this is exactly the coefficient of \(x^n\) in \(F(x)G(x)\). So “convolution”, “polynomial multiplication”, and “what NTT accelerates” are three names for one thing. The algorithm is the standard evaluate–multiply–interpolate:

1. evaluate \(F\) and \(G\) at the \(n\) powers of \(\omega\) (forward transform),
2. multiply pointwise — \(n\) multiplications instead of \(n^2\), because evaluation turns products of polynomials into products of numbers,
3. interpolate back (inverse transform, which is the same butterfly with \(\omega^{-1}\), scaled by \(n^{-1}\)).

The practical consequence: your entire job on a new problem is rewriting whatever you’re computing into the shape \(\sum_{i+j=n}\). If you can do that, the template does the rest. Everything below is about that rewrite.

One boundary condition worth stating: the transform is inherently cyclic of length \(n\). Indices that overflow past \(n\) wrap around and corrupt low coefficients. To get an honest linear convolution you must pad both inputs so \(n \geq \operatorname{len}(f) + \operatorname{len}(g) - 1\). Sometimes wraparound is what you want — see the next section.

Correlation is convolution with one side reversed

Problems rarely hand you \(\sum_{i+j=n}\). Far more often you get a correlation, where the indices subtract instead of adding:

“For every offset \(k\), slide \(g\) over \(f\) and take the overlap.” That’s the shape of anything asking “for every shift / every difference / every distance, count matches.”

Reversing one side converts it. Let \(f'[i] = f[K - 1 - i]\) for \(f\) of length \(K\). Then

substituting \(a = K - 1 - i\). The right side is \(\mathrm{corr}[n - K + 1]\), so:

Reverse, convolve, read from offset \(K-1\). That’s the whole bridge.

The cyclic case. When indices live mod \(K\) — rotations of a circular array, shifts of an alphabet — you want

You can fold the linear result, but the off-by-one bookkeeping is a reliable source of wrong answers. The cheaper move is to double \(g\): let \(G = g \mathbin\Vert g\), length \(2K\). Now \(G[i+k] = g[(i+k) \bmod K]\) for all \(i, k < K\) by construction, no wraparound reasoning required, and the linear formula applies verbatim.

The template

Iterative, in-place, decimation-in-time. conv pads to a power of two large enough to avoid wraparound; cyclic_correlate is the doubling trick above.

MOD = 998244353
G = 3  # primitive root of MOD

def ntt(a, invert):
    n = len(a)
    j = 0
    for i in range(1, n):                    # bit-reversal permutation
        bit = n >> 1
        while j & bit:
            j ^= bit
            bit >>= 1
        j |= bit
        if i < j:
            a[i], a[j] = a[j], a[i]
    length = 2
    while length <= n:
        w = pow(G, (MOD - 1) // length, MOD)  # primitive length-th root
        if invert:
            w = pow(w, MOD - 2, MOD)
        half = length >> 1
        for i in range(0, n, length):
            wn = 1
            for k in range(i, i + half):
                u = a[k]
                v = a[k + half] * wn % MOD
                a[k] = (u + v) % MOD
                a[k + half] = (u - v) % MOD
                wn = wn * w % MOD
        length <<= 1
    if invert:
        ninv = pow(n, MOD - 2, MOD)
        for i in range(n):
            a[i] = a[i] * ninv % MOD

def conv(f, g):
    need = len(f) + len(g) - 1
    n = 1
    while n < need:
        n <<= 1
    f = f + [0] * (n - len(f))
    g = g + [0] * (n - len(g))
    ntt(f, False)
    ntt(g, False)
    for i in range(n):
        f[i] = f[i] * g[i] % MOD
    ntt(f, True)
    return f[:need]

def cyclic_correlate(f, g, K):
    """c[k] = sum_i f[i] * g[(i + k) % K], for k = 0..K-1"""
    r = conv(f[::-1], g + g)
    return [r[k + K - 1] for k in range(K)]

The pow(G, (MOD-1)//length, MOD) line is the whole difference from an FFT: where a complex FFT would compute \(e^{2\pi i/\text{length}}\), this raises the field’s primitive root to the cofactor to land on an element of order exactly length. Note it needs length | MOD - 1, which is why n is rounded to a power of two and why the modulus was chosen with \(2^{23}\) in it.

ABC471 G — Caeser Syllables

ABC471 G — Caeser Syllables (AtCoder Beginner Contest 471), 600 points.

An alphabet of \(K\) symbols, each flagged vowel or not by \(V[0..K-1]\). Given \(A\) of length \(N\), for every shift \(k = 0 \dots K-1\) report the number of syllables of \(A'_i = (A_i + k) \bmod K\) — that is, the number of maximal runs of vowels.

\(N \leq 7 \times 10^6\), \(K \leq 2300\), 5 seconds.

\(A\) itself is generated by a PCG-style recurrence given as pseudocode, purely so the input file stays small. One trap there: the state advances only on the generated branch — for \(i \leq M\), where \(A_i\) is read from \(b_i\), the state is left alone. Advancing it unconditionally reproduces sample 2 shifted by two positions, which is a fun half hour to lose.

Step 1: kill the runs

Counting maximal runs directly is awkward because it’s a statement about boundaries. The standard rewrite: a run of vowels is a path, and for any set of paths, \(\text{components} = \text{vertices} - \text{edges}\). Here vertices are vowel positions and edges are adjacent vowel–vowel pairs:

Check on V V C V: 3 vowels, 1 adjacent vowel-vowel pair, 2 runs. This turns a boundary condition into two independent counts, each a plain sum with no “look at the neighbour and negate” logic.

Step 2: the first term is one correlation

Let \(\mathrm{cnt}[j]\) be how many times value \(j\) appears in \(A\) — one \(O(N)\) pass. Then

which is literally the cyclic correlation from the section above. One call, \(O(K \log K)\), all \(K\) shifts at once.

Step 3: the second term, and a claim that looks true but isn’t

Let \(c[x][y]\) count adjacent pairs \((A_{i-1}, A_i) = (x, y)\) — again one \(O(N)\) pass, into a \(K \times K\) table (5.3M entries at \(K = 2300\), fine). Then

It is very tempting at this point to argue: both symbols shift together, their difference \(d = (y - x) \bmod K\) never changes, so the pair’s contribution depends only on \(d\). That is false, and it’s worth being precise about why, because the true statement is right next to it. What depends only on \(d\) is the contribution summed over all \(k\). For a fixed \(k\) the term is \(V[u] \cdot V[(u+d) \bmod K]\) where \(u = (x+k) \bmod K\) — it depends on \(d\) and on where the pair currently sits, \(u\). Collapsing to \(d\) alone answers a question nobody asked (“over all rotations, how many transitions total”), and it’ll match on symmetric samples long enough to be convincing.

The correct move is to keep \(x\) and factor the sum instead:

The inner bracket is the same cyclic correlation as step 2, run once per row \(x\) — correlate row \(c[x][\cdot]\) against \(V\). That’s \(K\) correlations of length \(K\), then an \(O(K^2)\) combine. Note both operands’ roles: \(V\) is shared across every row, so its forward transform can be computed once and reused, halving the transform count.

\(O(N + K^2 \log K)\) overall. At \(N = 7 \times 10^6\) and \(K = 2300\) the \(O(N)\) generator pass is the part you actually have to micro-optimize; the NTT part is about \(7 \times 10^7\) modular operations.

Why one prime suffices

Every quantity here is a count of positions in \(A\), so it’s bounded by \(N \leq 7 \times 10^6 < 998244353\). No CRT, no three-prime dance — the mod never wraps, and the residues coming out of the inverse transform are the answers. This is the middle case from the first section, and recognizing it saves two thirds of the work.

Reference implementation

Using the template above. This is the direct transcription of the derivation, not a tuned submission — it’s what I checked the math against:

def solve(a, v, K):
    cnt = [0] * K
    for x in a:
        cnt[x] += 1
    term1 = cyclic_correlate(cnt, v, K)

    pair = [[0] * K for _ in range(K)]
    for i in range(1, len(a)):
        pair[a[i - 1]][a[i]] += 1

    term2 = [0] * K
    for x in range(K):
        if not any(pair[x]):
            continue
        Rx = cyclic_correlate(pair[x], v, K)
        for k in range(K):
            if v[(x + k) % K]:
                term2[k] += Rx[k]

    return [term1[k] - term2[k] for k in range(K)]

Verified against samples 1 and 2, and against a brute force over 300 random cases with \(K \leq 9\), \(N \leq 30\) — including, deliberately, asymmetric vowel masks, which is where the “depends only on \(d\)” shortcut breaks.

Takeaway

NTT is not here because \(N\) is huge. The \(O(N)\) scan over 7 million symbols is already optimal and no transform touches it. NTT enters only after the problem has been squeezed down to size \(K\), at the moment the question becomes “for every shift at once,” which is always a correlation, which is always a convolution. The reduction is the problem; the transform is a library call.