Introduzione ad FFT

Un’introduzione a FFT perché mi annoiavo

Il problema

Dati due polinomi P(x) e Q(x), vogliamo calcolarne il prodotto ans(x)=P(x)Q(x).

Un po’ di notazione

Un polinomio è univocamente determinato dai suoi coefficienti. Possiamo quindi rappresentare un polinomio A(x) di grado g con un array A di g + 1 elementi, tale che A[i] sia il coefficiente del monomio di grado i in A(x).
Per esempio se A(x) = 4 + 3x + 2x^2 - 3x^3, allora A = [4, 3, 2, -3].

Notiamo che \text{len}(A)\ge1+\deg(A). In generale la lunghezza della lista dei coefficienti è più comoda da usare rispetto al grado del polinomio.

Per questo d’ora in poi useremo N = \text{len}(P) e M = \text{len}(Q)

La soluzione naive

Con questa notazione possiamo descrivere agevolmente una soluzione naive del problema. È facile verificare che i coefficienti di ans(x)=P(x)Q(x) sono dati dalla formula ans[k]=\sum_{i+j=k}{P[i]Q[j]}.

Codice C++
#include <vector>
using namespace std;

vector<int64_t> prod(vector<int> P, vector<int> Q) {
    vector<int64_t> ans(P.size() + Q.size() - 1);

    for (int i = 0; i < P.size(); i++)
        for (int j = 0; j < Q.size(); j++)
            ans[i + j] += P[i] * Q[j];
    
    return ans;
}

Questa soluzione ha complessità \mathcal{O}(NM); si può fare di meglio, ma questa interpretazione è utile per ricondurre altri problemi alla moltiplicazione tra due polinomi.

La soluzione efficiente

Per trovare una soluzione più efficiente bisogna cambiare completamente punto di vista, sfruttando le proprietà dei polinomi. In particolare, è noto che dati n punti esiste un solo polinomio di grado (al più) n - 1 che li interpola (cioè che passa per tuti gli n punti). Per risolvere il problema allora è possibile usare il seguente algoritmo:

  1. Valutiamo P(x) e Q(x) su almeno N + M - 1 punti, cioè scelti k \ge N + M - 1 punti x_0, x_1, \dots, x_{k-1}, calcoliamo P(x_0), P(x_1), \dots, P(x_{k-1}) e Q(x_0), Q(x_1), \dots, Q(x_{k-1})

  2. Calcoliamo i valori ans(x_0) = P(x_0)Q(x_0), ans(x_1) = P(x_1)Q(x_1), \dots, ans(x_{k-1})=P(x_{k-1})Q(x_{k-1})

  3. Interpoliamo i valori trovati per ans(x_0), ans(x_1), \dots, ans(x_{k-1}) per trovare i coefficienti di ans.

Non sembra di aver fatto molti progressi: il secondo passaggio è chiaramente fattibile in \mathcal{O}(N + M), ma per il primo e il terzo non è evidente che esista una soluzione migliore di \mathcal{O}\left((N+M)^2\right).

Tuttavia un’accurata scelta dei valori di x_0, x_1, \dots, x_{k-1} porta a una soluzione più efficiente.

Passaggio 1

D’ora in poi lavoreremo con i numeri complessi e come valori per x_0, x_1, \dots, x_{k-1} sceglieremo le radici dell’unità.

Costruiamo una procedura ricorsiva FFT che prende in input un polinomio A(x) tale che \text{len}(A)=2^t per qualche t naturale e restituisce il valore di A(x) su tutte le radici 2^t-esime dell’unità.

Per fare ciò, costruiamo due polinomi A' = [A[0], A[2], \dots, A[2^t-2]] e A''=[A[1], A[3], \dots, A[2^t-1]].

Per esempio se A(x) = 4 + 3x + 2x^2 - 3x^3. allora A'(x) = 4 + 2x e A''(x)=3 - 3x.

Notiamo che \text{len}(A') = \text{len}(A'') = 2^{t-1} è ancora una potenza di 2, quindi possiamo chiamare ricorsivamente FFT su A' e su A''.

Notiamo inoltre che A(x)=A'(x^2)+xA''(x^2). Tuttavia se x è una radice 2^t-esima dell’unità, allora x^2 è una radice 2^{t-1}-esima dell’unità quindi tutti i valori necessari di A'(x) e A''(x) ci vengono già forniti dalle chiamate ricorsive e possiamo costruire i valori assunti da A(x) sulle radici dell’unità in \mathcal{O}(2^t)

Nel caso base della ricorsione, t=0 e stiamo quindi valutando un polinomio costante in 1. Possiamo quindi semplicemente restituire il termine noto del polinomio in input.

Complessità

Poiché T(n)=2T(n/2)+\mathcal{O}(n) possiamo concludere per il teorema dell’esperto che la procedura ricorsiva ha complessità \mathcal{O}(n \log n).

Passaggio 3

Per il passaggio 3 vogliamo ricostruire un polinomio di grado al più 2^t-1 sapendo i valori che il polinomio deve assumere sulle radici 2^t-esime dell’unità.

Di nuovo costruiamo una procedura ricorsiva IFFT che prende in input i valori assunti da un polinomio sulle radici 2^t-esime dell’unità e restituisce un polinomio A(x), tale che \text{len}(A) = 2^t che interpola i valori in input.

Chiamiamo V l’array di valori in input. In particolare V[j] è il valore che il polinomio restituito deve assumere in e^{\frac{2ij\pi}{2^t}}.

Costruiamo, similmente a prima, V'=[V[0], V[2], \dots, V[2^t-2]] e V''=[V[1], V[3], \dots, V[2^t-1]]. Per le stesse osservazioni di prima possiamo chiamare ricorsivamente IFFT su V' e V''. Chiamamo rispettivamente A'(x) e A''(x) i valori restituiti da queste due chiamate.

Claim: posto \omega = e^{\frac{2i\pi}{2^t}} vale
A[j] = \frac{1}{2}\left(A'[j]+\omega^{-j}A''[j]\right) per 0 \le j < 2^{t-1}
A[j] = \frac{1}{2}\left(A'[j -2^{t-1}]+\omega^{-j}A''[j -2^{t-1}]\right) per 2^{t-1} \le j < 2^t

Proof: Espandendo A(\omega^q) otteniamo:
\begin{align} & A(\omega^q) \\ &= \frac{1}{2}\sum_{j=0}^{2^{t-1}-1}{\left(A'[j] + \omega^{-j}A''[j]\right)\omega^{jq}} + \frac{1}{2}\sum_{j=2^{t-1}}^{2^t-1}{\left(A'[j - 2^{t-1}] + \omega^{-j}A''[j - 2^{t-1}]\right)\omega^{jq}} \\ &= \frac{1}{2}\sum_{j=0}^{2^{t-1}-1}{\left(A'[j] + \omega^{-j}A''[j]\right)\omega^{jq}} +\frac{1}{2} \sum_{j=0}^{2^{t-1}-1}{\left(A'[j] - \omega^{-j}A''[j]\right)(-\omega^j)^q} \\ &= \frac{1}{2}\sum_{j=0}^{2^{t-1}-1}{\left[\left(A'[j] + \omega^{-j}A''[j]\right)\omega^{jq} + \left(A'[j] - \omega^{-j}A''[j]\right)(-\omega^j)^q\right]} \\ &= \frac{1}{2}\sum_{j=0}^{2^{t-2}-1}{\left[\left(A'[2j] + \omega^{-2j}A''[2j]\right)\omega^{2jq} +\\ \left(A'[2j] - \omega^{-2j}A''[2j]\right)(-\omega^{2j})^q +\\ \left(A'[2j + 1] + \omega^{-2j - 1}A''[2j + 1]\right)\omega^{(2j + 1)q} +\\ \left(A'[2j + 1] - \omega^{-2j - 1}A''[2j + 1]\right)(-\omega^{2j + 1})^q\right]} \\ \end{align}

Ora se q è pari si ottiene:
\begin{align} &= \frac{1}{2}\sum_{j=0}^{2^{t-2}-1}{\left[\left(A'[2j] + \omega^{-2j}A''[2j]\right)\omega^{2jq} +\\ \left(A'[2j] - \omega^{-2j}A''[2j]\right)\omega^{2jq} + \\ \left(A'[2j + 1] + \omega^{-2j - 1}A''[2j + 1]\right)\omega^{(2j + 1)q} +\\ \left(A'[2j + 1] - \omega^{-2j - 1}A''[2j + 1]\right)\omega^{(2j + 1)q }\right]} \\ &= \frac{1}{2}\sum_{j=0}^{2^{t-2}-1}{\left[2\omega^{2jq}A'[2j] + 2\omega^{(2j + 1)q }A'[2j+1]\right]} \\ &= \sum_{j=0}^{2^{t-2}-1}{\left[\omega^{2jq}A'[2j] + \omega^{(2j + 1)q }A'[2j+1]\right]} \\ &= \sum_{j=0}^{2^{t-1}-1}{\omega^{jq}A'[j]} = \sum_{j=0}^{2^{t-1}-1}{(\omega^{q})^jA'[j]} = A'(\omega^q) \\ \end{align}

Se q è dispari invece:
\begin{align} &= \frac{1}{2}\sum_{j=0}^{2^{t-2}-1}{\left[\left(A'[2j] + \omega^{-2j}A''[2j]\right)\omega^{2jq} +\\ \left(A'[2j] - \omega^{-2j}A''[2j]\right)\left(-\omega^{2jq}\right) + \\ \left(A'[2j + 1] + \omega^{-2j - 1}A''[2j + 1]\right)\omega^{(2j + 1)q} +\\ \left(A'[2j + 1] - \omega^{-2j - 1}A''[2j + 1]\right)\left(-\omega^{(2j + 1)q }\right)\right]} \\ &= \frac{1}{2}\sum_{j=0}^{2^{t-2}-1}{\left[2\omega^{2jq}\omega^{-2j}A''[2j] + 2\omega^{(2j + 1)q}\omega^{-2j-1}A'[2j+1]\right]} \\ &= \sum_{j=0}^{2^{t-2}-1}{\left[\omega^{2j(q-1)}A''[2j] + \omega^{(2j + 1)(q-1)}A'[2j+1]\right]} \\ &= \sum_{j=0}^{2^{t-1}-1}{\omega^{j(q-1)}A''[j]} = \sum_{j=0}^{2^{t-1}-1}{(\omega^{q-1})^jA''[j]} = A''(\omega^{q-1}) \end{align}

In entrambi i casi la correttezza dell’interpolazione è garantita dalle condizioni sui valori restituiti dalle chiamate ricorsive.

Nel caso base della ricorsione, t=0 e dobbiamo quindi interpolare un singolo punto con un polinomio di grado 0. Possiamo quindi restituire il polinomio costante che passa per il punto dato.

Complessità

Come prima \mathcal{O}(n \log n). In realtà il passaggio 1 e il passaggio 3 sono molto più simili di quanto possa sembrare.

Implementazione

Con questo abbiamo finito di descrivere un algoritmo efficiente che moltiplica due polinomi. La complessità finale è \mathcal{O}((N+M)\log(N+M)).

Resta da chiarire come gestire il caso in cui \text{len}(P) o \text{len}(Q) non sono potenze di 2, tuttavia questo caso si gestisce semplicemente aggiungendo degli 0 in fondo alla lista dei coefficienti.

Bisogna tuttavia porre molta attenzione in come si implementa questo algoritmo, poiché implementare direttamente le procedure ricorsive descritte sopra (usando per esempio std::vector) produce del codice più lento della moltiplicazione naive anche per polinomi abbastanza grandi per via dell’overhead dovuto alle allocazioni e alle chiamate ricorsive. Per la maggior parte, le implementazioni efficienti di FFT sono iterative, tuttavia con qualche accortezza è possibile scrivere implementazioni ricorsive efficienti.
In particolare piuttosto che far restituire un array a ogni chiamata delle procedure ricorsive, è molto più efficiente modificare un array dichiarato esternamente, evitando in questo modo di dover fare lente allocazioni.
Nell’implementazione che segue la procedura fft sovrascrive un array che rappresenta un polinomio con uno che rappresenta il suo valore nelle radici dell’unità e ifft sovrascrive un array che rappresenta i valori nelle radici dell’unità con uno che rappresenta il polinomio che li interpola.

Codice C++
#include <bits/stdc++.h>
using namespace std;
 
#define MAXSIZE (1 << 19)
typedef complex<double> cx;
 
cx temp[MAXSIZE];
 
void fft(cx arr[], int step, int size) {
    if (size == 1) return;
 
    cx omega = exp(2i * (M_PI / size));
    fft(arr,        step * 2, size / 2);
    fft(arr + step, step * 2, size / 2);
 
    for (int j = 0; j < size; j++) temp[j] = arr[j * step];
 
    cx w = 1;
    for (int j = 0; j < size / 2; j++) {
        arr[j * step] =
            temp[j * 2] + w * temp[j * 2 + 1];
        arr[(j + size / 2) * step] =
            temp[j * 2] - w * temp[j * 2 + 1];
        
        w *= omega;
    }
}
 
void ifft(cx arr[], int step, int size) {
    if (size == 1) return;
 
    cx omega = exp(-2i * (M_PI / size));
    ifft(arr,        step * 2, size / 2);
    ifft(arr + step, step * 2, size / 2);
 
    for (int j = 0; j < size; j++) temp[j] = arr[j * step];
 
    cx w = 1;
    for (int j = 0; j < size / 2; j++) {
        arr[j * step] =
            temp[j * 2] + w * temp[j * 2 + 1];
        arr[(j + size / 2) * step] =
            temp[j * 2] - w * temp[j * 2 + 1];
        
        w *= omega;
    }
}
 
vector<int64_t> prod(vector<int> P, vector<int> Q) {
    int size = 1 << (int)ceil(log2(P.size() + Q.size() - 1));
    
    vector<cx> p(size), q(size);
 
    copy(P.begin(), P.end(), p.begin());
    copy(Q.begin(), Q.end(), q.begin());
 
    fft(p.data(), 1, size);
    fft(q.data(), 1, size);
    
    for (int i = 0; i < size; i++) p[i] *= q[i];
 
    ifft(p.data(), 1, size);
 
    vector<int64_t> ans(size);
    for (int i = 0; i < size; i++) ans[i] = p[i].real() / size + 0.5;
    return ans;
}

Limiti e varianti

Questo algoritmo lavora internamente con numeri a virgola mobile per rappresentare i complessi ed è pertanto soggetto a problemi di precisione.

Questo algoritmo può essere usato su qualunque campo che ammetta radici 2^t-esime dell’unità. Per esempio, oltre a \mathbb{C}, funziona anche sugli interi modulo p (per qualche p primo). Questa particolare istanza (chiamata NTT) permette di risolvere il problema del prodotto tra polinomi modulo p utilizzando solamente variabili intere e di conseguenza non ha problemi di precisione.

Problemi

5 Mi Piace

Grazie mille sig. Cervellesi per questa guida molto afftascinante :smiling_face_with_three_hearts:

2 Mi Piace

Altri problemi che usano FFT (e non sono “implementa FFT” ma non sono neanche impossibili):

Se volete, posso darvi hint (e su codeforces ci sono gli editorial).

3 Mi Piace