Proofs: Number-theoretic transform

NTT-MULTIPLY returns the product modulo p, exactly, in O(n log n) operations modulo p full

This proof shows: correctness, complexity.

Statement. Let $p$ be an odd prime, $g$ a primitive root modulo $p$, and $a, b$ integer sequences of lengths $\ell_a, \ell_b \ge 1$. Let $m = \ell_a + \ell_b - 1$ and let $n = 2^t$ be the least power of two with $n \ge m$, and suppose $n \mid p - 1$. Then NTT-MULTIPLY$(a, b, p, g)$ returns $c_k = \sum_{i + j = k} a_i b_j \bmod p$ for $0 \le k < m$. It does $3 n \lg n$ additions and subtractions and $\frac{3}{2} n \lg n + 2n$ multiplications modulo $p$, besides three tables of $\frac{n}{2} - 1$ multiplications each and $O(\log p)$ multiplications for the powers and inverses. If moreover $|c_k| < p/2$ for the integer coefficients, reading each residue in $(-p/2, p/2)$ gives them exactly; for $0 \le a_i < A$ and $0 \le b_j < B$ it suffices that $\min(\ell_a, \ell_b)(A - 1)(B - 1) < p$, reading residues in $[0, p)$.

All arithmetic below is in the field $F = \mathbb{Z}/p\mathbb{Z}$.

Step 1: the root. The group $F^\times$ is cyclic of order $p - 1$, generated by $g$. Line L4 sets $\omega = g^{(p-1)/n}$. If $\omega^d = 1$ with $0 < d$, then $p - 1$ divides $d (p-1)/n$, so $n \mid d$; and $\omega^n = g^{p-1} = 1$. So $\omega$ has order exactly $n$. If $t \ge 1$, then $z = \omega^{n/2} = g^{(p-1)/2}$ satisfies $z^2 = 1$ and $z \ne 1$ (the order of $g$ is $p - 1 > (p-1)/2$), and in a field the only roots of $X^2 - 1$ are $\pm 1$, so $z = -1$. The same holds for $\omega^{-1} = \omega^{n-1}$, which also has order $n$, and $(\omega^{-1})^{n/2} = (-1)^{-1} = -1$. Also, for $1 \le k \le n - 1$, $\omega^k \ne 1$ and $(\omega^k - 1)\sum_{j<n} \omega^{jk} = \omega^{nk} - 1 = 0$, so $\sum_{j<n} \omega^{jk} = 0$ in the field: $\omega$ is a principal $n$-th root of unity.

Step 2: the transform. Let $\zeta \in F$ with $\zeta^n = 1$ and $\zeta^{n/2} = -1$ if $n \ge 2$. We show that NTT$(x, \zeta, p)$ returns $\mathrm{DFT}_\zeta(x)$. Write $\mathrm{rev}_t$ for reversal of $t$ binary digits and $(Py)_i = y_{\mathrm{rev}_t(i)}$. Line L10 replaces $x$ by $Px$: it swaps each pair $\{j, \mathrm{rev}_t(j)\}$ with $j \ne \mathrm{rev}_t(j)$ once. Line L11 sets $W[k] = \zeta^k$, by induction on $k$. Let $T_t^\zeta$ be the map of lines L12 to L18. We prove $T_t^\zeta(Pz) = \mathrm{DFT}_\zeta(z)$ by induction on $t$; $t = 0$ is trivial.

For $t \ge 1$, the passes with $h \le n/4$ work on blocks of length $2h \le n/2$, so they never mix the two halves of the array. On each half, the pass with half length $h$ multiplies by $W[k \cdot n/(2h)] = (\zeta^2)^{k \cdot (n/2)/(2h)}$, which is the twiddle factor of the same pass for length $n/2$ and root $\zeta^2$; and $\zeta^2$ satisfies the hypothesis for $n/2$, since $(\zeta^2)^{n/4} = \zeta^{n/2} = -1$ when $n \ge 4$. For $q < n/2$ the top digit of $q$ is $0$, so $\mathrm{rev}_t(q) = 2\,\mathrm{rev}_{t-1}(q)$ and $\mathrm{rev}_t(n/2 + q) = 2\,\mathrm{rev}_{t-1}(q) + 1$. Hence the first half of $Pz$ is $P_{t-1} e$ with $e_r = z_{2r}$, and the second half is $P_{t-1} o$ with $o_r = z_{2r+1}$. By the induction hypothesis, after the first $t - 1$ passes the halves hold $E = \mathrm{DFT}_{\zeta^2}(e)$ and $O = \mathrm{DFT}_{\zeta^2}(o)$. The last pass ($h = n/2$, one block, stride $1$) maps $(E_k, O_k)$ for $k < n/2$ to $(E_k + \zeta^k O_k,\ E_k - \zeta^k O_k)$; each butterfly reads both cells before writing them, and different $k$ use disjoint cells. Splitting the DFT sum by the parity of $j$,

using $(\zeta^2)^{n/2} = 1$. For $k < n/2$ this is $E_k + \zeta^k O_k$, and for $k + n/2$ it is $E_k + \zeta^{k + n/2} O_k = E_k - \zeta^k O_k$. So the result is $\mathrm{DFT}_\zeta(z)$. Applying this with $z = x$ proves the claim.

Step 3: the convolution theorem. Let $A, B \in F^n$ and $(A \ast B)_r = \sum_{j<n} A_j B_{(r - j) \bmod n}$. Then $\mathrm{DFT}_\omega(A \ast B)_k = \sum_r \sum_j A_j B_{(r-j) \bmod n} \omega^{rk}$. Put $s = (r - j) \bmod n$; then $\omega^{rk} = \omega^{jk} \omega^{sk}$ since $\omega^n = 1$, and for fixed $j$ the map $r \mapsto s$ is a bijection of $\{0, \ldots, n-1\}$. So the sum is $\sum_j A_j \omega^{jk} \sum_s B_s \omega^{sk} = \hat A_k \hat B_k$.

Step 4: inversion. For $y \in F^n$, entry $j$ of $\mathrm{DFT}_{\omega^{-1}}(\mathrm{DFT}_\omega(y))$ is $\sum_r y_r \sum_k \omega^{(r - j)k}$. The inner sum is $n$ if $r = j$, and $0$ otherwise by Step 1 applied to $(r - j) \bmod n \in \{1, \ldots, n-1\}$. So it equals $n y_j$. Since $n < p$, $n$ is a unit of $F$.

Step 5: putting it together. By Steps 1 and 2, after line L6, $A = \mathrm{DFT}_\omega(a')$ and $B = \mathrm{DFT}_\omega(b')$ where $a', b'$ are the padded vectors of line L5. By Step 3, line L7 gives $C = \mathrm{DFT}_\omega(a' \ast b')$. By Step 2 with $\zeta = \omega^{-1}$ and Step 4, line L8 gives $n (a' \ast b')$, and line L9 gives $a' \ast b'$ in its first $m$ entries. Finally, $(a' \ast b')_k = \sum_{j} a'_j b'_{(k - j) \bmod n}$; a term is nonzero only if $j < \ell_a$ and $(k - j) \bmod n < \ell_b$. For $k < m \le n$ and $0 \le j < \ell_a$: if $j \le k$, then $(k - j) \bmod n = k - j$; if $j > k$, then $(k - j) \bmod n = n + k - j \ge n - \ell_a + 1 \ge \ell_b$, because $n \ge \ell_a + \ell_b - 1$, so that term is zero. Hence $(a' \ast b')_k = \sum_{i + j = k} a_i b_j \bmod p = c_k$.

Exact integers. The residue of $c_k$ determines $c_k$ uniquely among the integers in any interval of length $p$. If $|c_k| < p/2$, the representative in $(-p/2, p/2)$ is $c_k$. If $0 \le a_i < A$ and $0 \le b_j < B$, then each $c_k$ is a sum of at most $\min(\ell_a, \ell_b)$ products, each at most $(A-1)(B-1)$, so $0 \le c_k < p$ under the stated condition, and the representative in $[0, p)$ is $c_k$. For several primes $p_1, \ldots, p_r$, the Chinese remainder theorem gives a unique residue modulo $P = p_1 \cdots p_r$, and the same argument with $P$ in place of $p$ applies; Garner's form computes it as $c = r_1 + p_1 (y_2 + p_2 y_3)$ with $y_2 = (r_2 - r_1) p_1^{-1} \bmod p_2$ and $y_3 = (r_3 - r_1 - p_1 y_2)(p_1 p_2)^{-1} \bmod p_3$, which is congruent to $r_i$ modulo each $p_i$ by construction and lies in $[0, P)$.

Operation count. In NTT, each of the $t$ passes has $n/(2h) \cdot h = n/2$ butterflies, each with one multiplication, one addition and one subtraction modulo $p$ (line L16 and line L17), for $n \lg n$ additions and subtractions and $\frac{n}{2} \lg n$ multiplications in all, for every input; line L11 adds $\frac{n}{2} - 1$ multiplications. NTT-MULTIPLY calls NTT three times and adds $n$ multiplications in line L7 and $n$ in line L9, which gives the totals of the statement. The powers $g^{(p-1)/n}$, $\omega^{-1} = \omega^{n-1}$ and $n^{-1} = n^{p-2}$ cost $O(\log p)$ multiplications by repeated squaring. On the word RAM with $p^2 < 2^w$, each operation modulo $p$ is $O(1)$ word operations, and the index work is proportional to the number of butterflies plus $O(n)$, so the time is $\Theta(n \log n)$ for every input. The arrays $A$, $B$, $C$, the table and the reversed indices take $\Theta(n)$ words. $\blacksquare$