Fast Fourier Transform

import cmath, math

def fft(a, invert=False):
    """
    Cooley-Tukey FFT. a: coefficient list, length must be a power of 2.
    invert=False: DFT (coefficients -> samples at roots of unity)
    invert=True:  IDFT (samples -> coefficients)
    Returns list of complex numbers.
    O(n log n) time.
    """
    n = len(a)
    if n == 1:
        return list(a)                     # base case

    # Split even and odd coefficients
    a_even = fft(a[0::2], invert)         # a[0], a[2], a[4], ...
    a_odd  = fft(a[1::2], invert)         # a[1], a[3], a[5], ...

    # Primitive n-th root of unity
    angle = 2 * math.pi / n * (-1 if invert else 1)
    w  = complex(1)
    wn = cmath.exp(complex(0, angle))      # e^(i * 2pi/n) or conjugate

    result = [0] * n
    for k in range(n // 2):
        # Butterfly operation: the core of FFT
        result[k]         = a_even[k] + w * a_odd[k]
        result[k + n//2]  = a_even[k] - w * a_odd[k]
        w *= wn                            # advance to next root: w^(k+1)

    if invert:
        result = [x / n for x in result]  # IDFT: divide by n at end
    return result


def poly_multiply(a, b):
    """
    Multiply polynomials a and b using FFT. O(n log n).
    a, b: coefficient lists (constant term first).
    Returns coefficient list of product.
    """
    result_len = len(a) + len(b) - 1
    # Pad to next power of 2
    n = 1
    while n < result_len:
        n <<= 1
    fa = fft(a + [0] * (n - len(a)))      # FFT of a (zero-padded)
    fb = fft(b + [0] * (n - len(b)))      # FFT of b (zero-padded)
    # Pointwise multiply in sample representation: O(n)
    fc = [fa[k] * fb[k] for k in range(n)]
    # IFFT to get product coefficients
    c = fft(fc, invert=True)
    return [round(x.real) for x in c[:result_len]]

# Example: (1 + 2x + 3x^2) * (4 + 5x) = 4 + 13x + 22x^2 + 15x^3
print(poly_multiply([1, 2, 3], [4, 5]))   # [4, 13, 22, 15]

import cmath, math, time

def dft_naive(a):
    """Naive O(n^2) DFT for comparison."""
    n = len(a)
    result = []
    for k in range(n):
        val = sum(a[j] * cmath.exp(-2j * math.pi * k * j / n)
                  for j in range(n))
        result.append(val)
    return result

def next_pow2(n):
    p = 1
    while p < n: p <<= 1
    return p

def benchmark_fft_vs_naive(sizes=[32, 64, 128, 256, 512]):
    """
    Compare DFT O(n^2) vs FFT O(n log n) empirically.
    Shows speedup grows rapidly with n.
    """
    import random
    results = []
    for n in sizes:
        a = [random.random() for _ in range(n)]

        t0 = time.perf_counter()
        dft_naive(a)
        t_naive = time.perf_counter() - t0

        t0 = time.perf_counter()
        fft(a + [0] * (next_pow2(n) - n))
        t_fft = time.perf_counter() - t0

        results.append((n, t_naive, t_fft, t_naive/t_fft))
        print(f"n={n:4d}: naive={t_naive*1000:.2f}ms  "
              f"fft={t_fft*1000:.2f}ms  "
              f"speedup={t_naive/t_fft:.1f}x")
    return results

# Key identity: for n-th roots of unity omega_n^k = e^(2*pi*i*k/n)
# FFT evaluates A at {omega_n^0, omega_n^1, ..., omega_n^(n-1)}
# IFFT recovers A from those samples: multiply by conjugate roots, divide by n

def demo_round_trip():
    """Verify FFT followed by IFFT recovers original coefficients."""
    original = [1, 2, 3, 4, 5, 6, 7, 8]   # must be power-of-2 length
    samples  = fft(original)                # DFT: coefficients -> samples
    recovered = fft(samples, invert=True)   # IDFT: samples -> coefficients
    print([round(x.real) for x in recovered])  # [1, 2, 3, 4, 5, 6, 7, 8]

import numpy as np

# ── FAST CONVOLUTION (polynomial multiplication via NumPy FFT) ───────
def fast_convolve(a, b):
    """
    Convolve sequences a and b using FFT. O(n log n).
    Equivalent to np.convolve(a, b) but faster for large inputs.
    """
    n = len(a) + len(b) - 1
    # Pad to next power of 2 for efficiency
    N = 1 << (n - 1).bit_length()
    fa = np.fft.rfft(a, N)             # real FFT (2x faster for real input)
    fb = np.fft.rfft(b, N)
    fc = fa * fb                       # pointwise multiply: O(N)
    c  = np.fft.irfft(fc, N)          # IFFT
    return np.round(c[:n]).astype(int)

# ── LOW-PASS FILTER (zero out high frequencies) ─────────────────────
def low_pass_filter(signal, cutoff_fraction=0.1):
    """
    Keep only lowest cutoff_fraction of frequencies.
    Used in: MP3 compression, smoothing, noise reduction.
    """
    n    = len(signal)
    freq = np.fft.rfft(signal)
    cutoff = int(len(freq) * cutoff_fraction)
    freq[cutoff:] = 0                  # zero out high frequencies
    return np.fft.irfft(freq, n)

# ── LARGE INTEGER MULTIPLICATION ──────────────────────────────────────
def large_int_multiply(x, y):
    """
    Multiply large integers x, y by treating their digits as polynomials.
    poly_multiply([d0, d1, ...], [e0, e1, ...]) then handle carries.
    O(n log n) for n-digit numbers — beats naive O(n^2).
    """
    digits_x = [int(d) for d in str(x)][::-1]  # LSB first
    digits_y = [int(d) for d in str(y)][::-1]
    product_digits = fast_convolve(digits_x, digits_y)

    # Handle carries
    result, carry = [], 0
    for d in product_digits:
        total = d + carry
        result.append(total % 10)
        carry = total // 10
    while carry:
        result.append(carry % 10)
        carry //= 10
    return int(''.join(str(d) for d in result[::-1]))

# Example: 12345 * 67890
print(large_int_multiply(12345, 67890))   # 838102050