Python - Number Theoretic Transformation

The Number Theoretic Transformation (NTT) is a mathematical technique used for efficient polynomial multiplication and convolution operations. It's closely related to the Fast Fourier Transform (FFT) but operates in finite fields, making it particularly useful in cryptography, error-correcting codes, and number theory applications.

In this article, we'll explore two implementations of NTT using the Cooley-Tukey algorithm: one using complex numbers (traditional FFT) and another using modular arithmetic with integers.

Method 1: Using Complex Numbers (Traditional FFT)

The Cooley-Tukey FFT algorithm can be adapted for NTT by working with complex exponentials. This approach uses the divide-and-conquer strategy to recursively compute the transformation ?

import cmath

def fft_cooley_tukey(x):
    """Compute FFT using Cooley-Tukey algorithm with complex numbers"""
    N = len(x)
    
    # Base case
    if N == 1:
        return x
    
    # Divide: split into even and odd indexed elements
    even_part = fft_cooley_tukey(x[::2])
    odd_part = fft_cooley_tukey(x[1::2])
    
    # Conquer: compute twiddle factors and combine
    T = [cmath.exp(-2j * cmath.pi * k / N) * odd_part[k] for k in range(N // 2)]
    
    # Combine results
    result = ([even_part[k] + T[k] for k in range(N // 2)] + 
              [even_part[k] - T[k] for k in range(N // 2)])
    
    return result

# Example usage
input_array = [1, 2, 3, 4]
output = fft_cooley_tukey(input_array)

print("FFT using Complex Numbers:")
print("Input:", input_array)
print("Output:", [complex(round(x.real, 2), round(x.imag, 2)) for x in output])

FFT using Complex Numbers:
Input: [1, 2, 3, 4]
Output: [(10+0j), (-2+2j), (-2+0j), (-2-2j)]

Method 2: Using Modular Arithmetic (Integer NTT)

This approach performs NTT using modular arithmetic in a finite field. It's more suitable for applications requiring exact integer computations ?

import numpy as np

def ntt_cooley_tukey(x, p, g):
    """Compute NTT using Cooley-Tukey algorithm with modular arithmetic"""
    N = len(x)
    
    # Base case
    if N == 1:
        return x
    
    # Recursively compute even and odd parts
    even_part = ntt_cooley_tukey(x[::2], p, (g * g) % p)
    odd_part = ntt_cooley_tukey(x[1::2], p, (g * g) % p)
    
    # Initialize result array
    result = np.zeros(N, dtype=int)
    factor = 1
    
    # Combine results using modular arithmetic
    for i in range(N // 2):
        term = (factor * odd_part[i]) % p
        result[i] = (even_part[i] + term) % p
        result[i + N // 2] = (even_part[i] - term) % p
        factor = (factor * g * g) % p
    
    return result

# Example usage
input_array = np.array([1, 2, 3, 4])
prime = 17  # Using a larger prime for better demonstration
primitive_root = 3  # Primitive root for prime 17

output = ntt_cooley_tukey(input_array, prime, primitive_root)

print("NTT using Modular Arithmetic:")
print("Input:", input_array.tolist())
print("Prime modulus:", prime)
print("Primitive root:", primitive_root)
print("Output:", output.tolist())

NTT using Modular Arithmetic:
Input: [1, 2, 3, 4]
Prime modulus: 17
Primitive root: 3
Output: [10, 15, 6, 2]

Comparison

Key Applications

Complex FFT: Ideal for signal processing, image processing, and applications where approximate results are acceptable.

Integer NTT: Essential in cryptographic algorithms, polynomial arithmetic over finite fields, and when exact integer results are required.

Conclusion

Both NTT implementations serve different purposes: complex FFT for general polynomial operations and integer NTT for exact modular computations. Choose the complex method for signal processing applications and the modular method for cryptographic and number-theoretic computations where precision is critical.