2つの多項式の積を求めるプログラム

convolution_fft.py

import cmath


def fft(a, inv):
    n = len(a)
    if n == 1:
        return a
    even = []
    odd = []
    for i in range(n):
        if i % 2 == 0:
            even.append(a[i])
        else:
            odd.append(a[i])
    d_even = fft(even, inv)
    d_odd = fft(odd, inv)
    dn = float(n)
    zeta = cmath.rect(1.0, 2 * cmath.pi * inv / dn)
    now = 1
    ret = []
    for i in range(n):
        ret.append(d_even[i % (n // 2)] + now * d_odd[i % (n // 2)])
        now *= zeta
    return ret


def convolution(a, b):
    len_a = len(a)
    len_b = len(b)
    len_c = len_a + len_b - 1  # len_c は A(x) * B(x) の次数
    n = 1
    while n <= len_c:
        n *= 2

    # 配列の長さが n になるまで、配列の末尾に0を追加する
    while len(a) < n:
        a.append(0)
    while len(b) < n:
        b.append(0)

    # A(x) の FFT DA(t), B(x) の FFt DB(t) を求める。
    da = fft(a, 1)
    db = fft(b, 1)

    dc = []
    for i in range(n):
        dc.append(da[i] * db[i])

    # C(x) は DC(t) を IFFT すれば求まる。
    c = fft(dc, -1)

    # IFFT の後は最後に n で割ることを忘れずに。
    ret = []
    for i in range(n):
        ret.append(c[i].real / float(n))
    return ret


a = list(map(int, input().split()))
b = list(map(int, input().split()))
na = len(a)
nb = len(b)
c = convolution(a, b)
for i in range(na + nb - 1):
    print(round(c[i]))

FFTを使って2つの多項式の積を求めるプログラム - Mae向きなブログ