import math import cmath def dft(p): """Quadratic formula from the definition.""" N = len(p) w = cmath.exp(-2*math.pi*1j/N) # primitive N-th square root of 1 x = [0]*N y = [0]*N for k in range(N): x[k] = w**k for n in range(N): y[k] += p[n] * x[k]**n return y def compute_inverse(N1, N2): for i in range(N2): if N1*i % N2 == 1: return i raise ValueError("Failed to find an inverse to {0} mod {1}.".format(N1, N2)) def prime_fft(p, divisors): """https://en.wikipedia.org/wiki/Prime-factor_FFT_algorithm""" if len(divisors) == 1: return dft(p) N = len(p) N1 = divisors[0] N2 = N//N1 N1_inv = compute_inverse(N1, N2) N2_inv = compute_inverse(N2, N1) ys = [] for n1 in range(N1): # compute rows p_ = [p[(n2*N1+n1*N2) % N] for n2 in range(N2)] ys.append(prime_fft(p_, divisors[1:])) for k2 in range(N2): # compute cols p_ = [row[k2] for row in ys] y_ = dft(p_) for (yi, row) in zip(y_, ys): # update col row[k2] = yi # remap and output res = [0]*N for k1 in range(N1): for k2 in range(N2): res[(k1*N2*N2_inv+k2*N1*N1_inv) % N] = ys[k1][k2] return res # arr = [9,8,5,3,8,4,9,7,0,9,5,6,6,2,4] # print(dft(arr)) # print() # print(prime_fft(arr,[3,5])) """ def fft(x): N = len(x) if N <= 1: return x even = fft(x[0::2]) odd = fft(x[1::2]) T = [cmath.exp(-2j*math.pi*k/N)*odd[k] for k in range(N//2)] return [even[k] + T[k] for k in range(N//2)] + \ [even[k] - T[k] for k in range(N//2)] def fft_mixed(x, r1, r2): A = dft(x) for k0 in range(0, r2): pass def bit_reverse(x, width): y = 0 for i in range(width): y <<= 1 y |= x&1 x >>= 1 return y def bit_reverse_seq(x): n = len(x) width = round(math.log2(n)) return [x[bit_reverse(i, width)] for i in range(n)] def fft2(x): y = bit_reverse_seq(x) n = len(x) omegas = [(-2*math.pi*1j/n)**i for i in range(0, n)] b = 1 while b