다항식 기초: FFT와 NTT

이 문서는 다항식 연산의 기본 개념을 다룹니다. 특히 빠른 푸리에 변환(FFT)과 빠른 수론 변환(NTT)을 중심으로 설명합니다.

FFT: 다항식 곱셈의 효율적 구현

FFT는 다항식을 단위근에서의 점값 표현으로 변환하는 과정을 가속화하는 알고리즘입니다. 주어진 다항식 F(x) = Σ aᵢxⁱωₙ⁰, ωₙ¹, ..., ωₙ^(n-1)에서 평가하는 과정을 분할 정복 방식으로 수행합니다.

핵심 아이디어는 다음과 같습니다:

  • F(ωₙ^k) = F⁰(ωₙ/₂ᵏ) + ωₙᵏ × F¹(ωₙ/₂ᵏ)
  • 짝수 차수 항과 홀수 차수 항을 나누어 각각 재귀적으로 계산한 후 결합

이를 통해 시간 복잡도를 O(n log n)로 줄일 수 있습니다.

IDFT: 점값에서 계수 복원

역변환은 F'(ωₙ⁻ᵏ) = n × aₖ의 성질을 이용합니다. 이를 통해 점값 표현에서 원래 계수를 회복할 수 있습니다. 이때 단위근의 반전 및 정규화가 필요하며, 최종 결과에는 n으로 나누는 단계가 포함됩니다.

FFT 구현 예시

struct Complex {
    double real, imag;
    Complex(double r = 0, double i = 0) : real(r), imag(i) {}
    
    Complex operator+(const Complex& b) const {
        return {real + b.real, imag + b.imag};
    }
    
    Complex operator-(const Complex& b) const {
        return {real - b.real, imag - b.imag};
    }
    
    Complex operator*(const Complex& b) const {
        return {
            real * b.real - imag * b.imag,
            real * b.imag + imag * b.real
        };
    }
};

Complex A[N], B[N];
int limi = 1, len;
int rev[N];

void fft(Complex* a, int sign) {
    for (int i = 0; i < limi; ++i)
        if (i < rev[i])
            swap(a[i], a[rev[i]]);
    
    for (int step = 1; step < limi; step <<= 1) {
        double angle = sign * M_PI / step;
        Complex w(cos(angle), sin(angle));
        
        for (int block = 0; block < limi; block += (step << 1)) {
            Complex curr(1, 0);
            for (int i = 0; i < step; ++i, curr = curr * w) {
                Complex u = a[block + i];
                Complex v = curr * a[block + i + step];
                a[block + i] = u + v;
                a[block + i + step] = u - v;
            }
        }
    }
}

NTT: 정수 모듈러 환에서의 변환

NTT는 정수 범위 내에서 유사한 기능을 수행하지만, 복소수 대신 모듈러 연산을 사용합니다. 일반적으로 소수 P = 998244353와 그 원근 G = 3를 사용합니다.

NTT의 핵심은 원근 G를 이용해 단위근과 유사한 순환 성질을 구현하는 것입니다.

NTT 구현 예시

const int MOD = 998244353;
const int G = 3;
const int Gi = (MOD + 1) / G;

long long A[N], B[N];
int limi = 1, log_len;
int rev[N];

long long pow_mod(long long base, long long exp) {
    long long res = 1;
    while (exp) {
        if (exp & 1) res = res * base % MOD;
        base = base * base % MOD;
        exp >>= 1;
    }
    return res;
}

void ntt(long long* a, int type) {
    for (int i = 0; i < limi; ++i)
        if (i < rev[i])
            swap(a[i], a[rev[i]]);
    
    for (int step = 1; step < limi; step <<= 1) {
        long long w = pow_mod(type == 1 ? G : Gi, (MOD - 1) / (step << 1));
        for (int block = 0; block < limi; block += (step << 1)) {
            long long curr = 1;
            for (int i = 0; i < step; ++i, curr = curr * w % MOD) {
                long long u = a[block + i];
                long long v = curr * a[block + i + step] % MOD;
                a[block + i] = (u + v) % MOD;
                a[block + i + step] = (u - v + MOD) % MOD;
            }
        }
    }
    
    if (type == -1) {
        long long inv = pow_mod(limi, MOD - 2);
        for (int i = 0; i < limi; ++i)
            a[i] = a[i] * inv % MOD;
    }
}

응용: 고정밀 곱셈

큰 수의 곱셈은 다항식 곱셈과 동일한 구조를 가집니다. 각 자릿수를 계수로 보고, FFT 또는 NTT로 O(n log n) 시간에 처리할 수 있습니다.

최적화 기법

  • 원근 사전 계산: 각 단계별 원근 값을 미리 계산하여 루프 내에서의 연산을 줄임
  • 반전 배열 활용: 비트 역전 인덱싱을 통해 메모리 접근 패턴 개선
  • 나눗셈 최적화: 역원을 미리 계산하여 나누기 연산을 곱하기로 치환

예제 문제: 회전된 배열의 최대 내적

문제에서 회전 가능한 두 배열의 내적을 최대화해야 할 때, 이를 다항식 곱셈의 결과에서 특정 위치의 값들을 조사함으로써 해결합니다. 이 경우 NTT를 사용하여 O(n log n)으로 처리 가능합니다.

주의사항

  • 모든 연산은 반드시 모듈러 연산을 적용
  • 반전 인덱스는 i < rev[i] 조건에서 스왑
  • 첫 번째 루프는 i = 1부터 시작하고 i <<= 1로 증가
  • 세 번째 루프의 범위는 p < step이며, t *= w는 반드시 모듈러 연산 포함
  • 역변환은 type == -1 시에만 수행

연습 문제

태그: FFT NTT 다항식 곱셈 수론 변환 고정밀 곱셈

7월 31일 20:45에 게시됨