高速フーリエ変換 (fast Fourier transform)
(NumberTheory/Convolution/fft_convolve.py)

• View this file on GitHub
• Last update: 2021-06-13 01:12:48+09:00

使い方

fft_convolve(a: Sequence[int], b: Sequence[int]) -> List[int]
長さ $N$ の数列 $a$ と長さ $M$ の数列 $b$ について、$c_k = \sum_{i + j \equiv k} a_ib_j$ となる長さ $N + M - 1$ の数列 $c$ を返す。計算量 $O((N + M) \log (N + M))$

Verified with

• TestCase/yukicoder/yuki0206.test.py

Code

from cmath import pi, exp


def _fft(a, h):
    roots = [exp(2.0j * pi / 2 ** i) for i in range(h + 1)]
    for i in range(h):
        m = 1 << (h - i - 1)
        for j in range(1 << i):
            w = 1
            j *= 2 * m
            for k in range(m):
                a[j + k], a[j + k + m] = \
                    a[j + k] + a[j + k + m], (a[j + k] - a[j + k + m]) * w
                w *= roots[h - i]


def _ifft(a, h):
    iroots = [exp(-2.0j * pi / 2 ** i) for i in range(h + 1)]
    for i in range(h):
        m = 1 << i
        for j in range(1 << (h - i - 1)):
            w = 1
            j *= 2 * m
            for k in range(m):
                a[j + k], a[j + k + m] = \
                    a[j + k] + a[j + k + m] * w, a[j + k] - a[j + k + m] * w
                w *= iroots[i + 1]
    n = 1 << h
    for i in range(n):
        a[i] /= n


def fft_convolve(a, b):
    len_ab = len(a) + len(b)
    n = 1 << (len(a) + len(b) - 1).bit_length()
    h = n.bit_length() - 1
    a = list(a) + [0] * (n - len(a))
    b = list(b) + [0] * (n - len(b))

    _fft(a, h), _fft(b, h)
    a = [va * vb for va, vb in zip(a, b)]
    _ifft(a, h)
    return [int(abs(a[i]) + 0.5) for i in range(len_ab - 1)]

Traceback (most recent call last):
  File "/opt/hostedtoolcache/Python/3.12.4/x64/lib/python3.12/site-packages/onlinejudge_verify/documentation/build.py", line 71, in _render_source_code_stat
    bundled_code = language.bundle(stat.path, basedir=basedir, options={'include_paths': [basedir]}).decode()
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/hostedtoolcache/Python/3.12.4/x64/lib/python3.12/site-packages/onlinejudge_verify/languages/python.py", line 96, in bundle
    raise NotImplementedError
NotImplementedError

Back to top page