『Algorithms in C アルゴリズム』の「41 高速フーリエ変換」

41_FFT.py

import cmath


def eval(p, N, k):
    if N == 2:
        p0, p1 = p[k], p[k + 1]
        p[k], p[k + 1] = p0 + p1, p0 - p1
    else:
        for i in range(N // 2):
            j = k + 2 * i
            t[i] = p[j]
            t[i + N // 2] = p[j + 1]
        for i in range(N):
            p[k + i] = t[i]
        eval(p, N // 2, k)
        eval(p, N // 2, k + N // 2)
        j = outN // N
        for i in range(N // 2):
            p0 = w[i * j] * p[k + (N // 2) + i]
            t[i] = p[k + i] + p0
            t[i + N // 2] = p[k + i] - p0
        for i in range(N):
            p[k + i] = t[i]


N = 4  # 2の冪乗にする必要がある。
outN = 2 * N
t = [0] * outN
w = [cmath.rect(1.0, 2 * cmath.pi * j / outN) for j in range(outN)]

p = [1, 1, 1, 0, 0, 0, 0, 0]  # p(x) = 1 + x + x^2
q = [2, -1, 1, 0, 0, 0, 0, 0]  # q(x) = 2 - x + x^2
eval(p, outN, 0)
eval(q, outN, 0)
r = []
for i in range(outN):
    r.append(p[i] * q[i])
eval(r, outN, 0)
for i in range(1, N):
    r[i], r[outN - i] = r[outN - i], r[i]
for i in range(outN):
    r[i] = r[i].real / float(outN)
for i in range(N + 1):
    print(round(r[i]))

『アルゴリズムC 第3巻』の「41 高速フーリエ変換」をPythonで - Mae向きなブログ