Number theoretic transform

ntt.py

# Number theoretic transform.
# Not O(n log n) but O(n^2).

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

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

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

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

def NTT(L, w, inv=False):
    n = len(L)
    ret = [reduce(lambda x, y: (x+y) % P, [mulmod(pow(w, (i*j), P), L[j]) for j in range(n)]) for i in range(n)]
    if not inv:
        return ret
    ni = -(3<<(30-pow2(n)))
    return [mulmod(e, ni) for e in ret]

def test():
    w, W = omega(8)
    a = [1, 2, 1, 0, 0, 0, 0, 0]
    b = [1, 1, 0, 0, 0, 0, 0, 0]
    A = NTT(a, w)
    B = NTT(b, w)
    C = [mulmod(p, q) for (p, q) in zip(A, B)]
    c = NTT(C, W, inv=True)
    print c # [1, 3, 3, 1, 0, 0, 0, 0]

test()