Fast Number theoretic transform

FNTT.py

# number theoretic transform implemented with FFT algorithm: O(n log n).

P = (3<<30)+1 # 3221225473
omega0 = 5
omega0inv = 1932735284

def mulmod(a, b):
    return (a*b) % P

OMEGA = {}

def omega(n):
    e = (P-1)/n
    return pow(omega0, e, P), pow(omega0inv, e, P)

for i in range(31):
    OMEGA[1<<i] = omega(1<<i)

def pow2(n):
    if n == 1:
        return 0
    if n & 1:
        return -1
    return pow2(n>>1) + 1

def diag_trf(L, inv):
    n = len(L)<<1
    w = OMEGA[n][inv]
    return [mulmod(e, pow(w, i, P)) for (i, e) in enumerate(L)]

def FNTT(L, inv=0):
    n = len(L)
    if n <= 1:
        return L
    factor = -(3<<(30-pow2(n)))
    if inv != -1:
        factor = 1
    elif inv == -1:
        inv = 1
    former = FNTT(L[::2], inv)
    latter = diag_trf(FNTT(L[1::2], inv), inv)
    ret = [(p + q) % P for (p, q) in zip(former, latter)] + [(p - q) % P for (p, q) in zip(former, latter)]
    return [mulmod(e, factor) for e in ret]

def test():
    a = [1, 2, 1, 0, 0, 0, 0, 0]
    b = [1, 1, 0, 0, 0, 0, 0, 0]
    A = FNTT(a)
    B = FNTT(b)
    C = [mulmod(p, q) for (p, q) in zip(A, B)]
    c = FNTT(C, inv=-1)
    print c # [1, 3, 3, 1, 0, 0, 0, 0]

test()