[수학] FFT와 NTT (Fast Fourier Transform & Number Theoretic Transform)

Date:     Updated:

카테고리:

태그:

💡 FFT와 NTT 개요

FFT(Fast Fourier Transform)와 NTT(Number Theoretic Transform)는 다항식 곱셈을 빠르게 수행하기 위한 알고리즘이다.

두 다항식의 곱셈을 naive하게 구현하면 $O(N^2)$의 시간 복잡도가 필요하지만, FFT나 NTT를 사용하면 $O(N \log N)$의 시간 복잡도로 해결할 수 있다.

🎯 알고리즘의 목적

  • 큰 수의 곱셈을 빠르게 수행
  • 다항식 곱셈의 효율적인 계산
  • 신호 처리 및 데이터 분석
  • 문자열 매칭 최적화

📚 FFT (Fast Fourier Transform)

1. FFT의 기본 원리

FFT는 다항식을 계수 표현(coefficient representation)에서 점값 표현(point-value representation)으로 변환하는 알고리즘이다.

다항식 $A(x) = a_0 + a_1x + a_2x^2 + … + a_{n-1}x^{n-1}$가 있을 때:

  1. 계수 표현: $[a_0, a_1, a_2, …, a_{n-1}]$
  2. 점값 표현: $[(x_0, y_0), (x_1, y_1), …, (x_{n-1}, y_{n-1})]$ where $y_i = A(x_i)$

2. FFT의 핵심 아이디어

FFT는 복소수 단위근(complex roots of unity)을 사용한다. n차 단위근은 다음과 같다:

$\omega_n = e^{2\pi i/n}$

주요 성질:

  • $\omega_n^n = 1$
  • $\omega_n^{k+n/2} = -\omega_n^k$

3. FFT 구현의 핵심 단계

void fft(vector<complex<double>>& a, bool invert) {
    int n = a.size();
    if (n == 1) return;

    vector<complex<double>> a0(n/2), a1(n/2);
    for (int i = 0; 2 * i < n; i++) {
        a0[i] = a[2*i];
        a1[i] = a[2*i+1];
    }
    
    fft(a0, invert);
    fft(a1, invert);

    double ang = 2 * PI / n * (invert ? -1 : 1);
    complex<double> w(1), wn(cos(ang), sin(ang));
    
    for (int i = 0; 2 * i < n; i++) {
        a[i] = a0[i] + w * a1[i];
        a[i + n/2] = a0[i] - w * a1[i];
        if (invert) {
            a[i] /= 2;
            a[i + n/2] /= 2;
        }
        w *= wn;
    }
}

📚 NTT (Number Theoretic Transform)

1. NTT의 특징

NTT는 FFT의 변형으로, 복소수 대신 모듈러 연산을 사용한다. 이는 다음과 같은 장점이 있다:

  • 정수만 사용하므로 부동소수점 오차가 없음
  • 모듈러 연산으로 오버플로우 방지
  • 더 빠른 연산 속도

2. NTT를 위한 모듈러와 원시근

NTT를 사용하기 위해서는 다음 조건을 만족하는 소수 $P$와 원시근 $g$가 필요하다:

  • $P = c \cdot 2^k + 1$ 형태의 소수 (c는 작은 홀수)
  • $g$는 $P$의 원시근

자주 사용되는 값들:

  • $P = 998244353 = 119 \cdot 2^{23} + 1$
  • $g = 3$

3. NTT 구현

const int MOD = 998244353;
const int ROOT = 3;

int power(int a, int n) {
    int ret = 1;
    while (n) {
        if (n & 1) ret = (1LL * ret * a) % MOD;
        a = (1LL * a * a) % MOD;
        n >>= 1;
    }
    return ret;
}

void ntt(vector<int>& a, bool invert) {
    int n = a.size();
    for (int i = 1, j = 0; i < n; i++) {
        int bit = n >> 1;
        for (; j & bit; bit >>= 1)
            j ^= bit;
        j ^= bit;
        if (i < j)
            swap(a[i], a[j]);
    }

    for (int len = 2; len <= n; len <<= 1) {
        int wlen = invert ? power(ROOT, MOD - 1 - (MOD - 1) / len)
                         : power(ROOT, (MOD - 1) / len);
        for (int i = 0; i < n; i += len) {
            int w = 1;
            for (int j = 0; j < len / 2; j++) {
                int u = a[i + j];
                int v = (1LL * a[i + j + len / 2] * w) % MOD;
                a[i + j] = (u + v) % MOD;
                a[i + j + len / 2] = (u - v + MOD) % MOD;
                w = (1LL * w * wlen) % MOD;
            }
        }
    }

    if (invert) {
        int n_inv = power(n, MOD - 2);
        for (int& x : a)
            x = (1LL * x * n_inv) % MOD;
    }
}

🎨 응용 예제

1. 큰 수의 곱셈

두 큰 수 A와 B의 곱셈을 수행할 때:

  1. 각 수를 다항식으로 변환
  2. FFT/NTT를 사용해 다항식 곱셈
  3. 결과를 다시 수로 변환
vector<int> multiply(vector<int>& a, vector<int>& b) {
    int n = 1;
    while (n < a.size() + b.size())
        n <<= 1;
    a.resize(n);
    b.resize(n);
    
    ntt(a, false);
    ntt(b, false);
    
    for (int i = 0; i < n; i++)
        a[i] = (1LL * a[i] * b[i]) % MOD;
        
    ntt(a, true);
    return a;
}

2. 문자열 매칭

문자열 매칭에서 FFT/NTT를 활용하면 와일드카드가 포함된 패턴 매칭을 효율적으로 처리할 수 있다.

❓ 주의사항과 팁

  1. 정밀도 문제
    • FFT는 부동소수점 연산을 사용하므로 정밀도 오차가 발생할 수 있다.
    • 정밀도가 중요한 경우 NTT를 사용하는 것이 좋다.
  2. 메모리 사용량
    • 입력 크기의 2의 거듭제곱까지 패딩이 필요하다.
    • 큰 입력의 경우 메모리 사용량에 주의가 필요하다.
  3. 최적화 팁
    • 작은 크기의 입력에는 일반적인 곱셈이 더 빠를 수 있다.
    • 병렬화가 가능한 구조이므로 필요시 병렬 처리를 고려한다.

🔍 참고 자료

Mathematics Theory 카테고리 내 다른 글 보러가기