Transformada Rápida de Fourier (FFT)

En este artículo discutiremos un algoritmo que permite multiplicar dos polinomios de longitud $n$ en tiempo $O(n \log n)$ , lo cual es mejor que la multiplicación trivial que toma tiempo $O(n^2)$ . Obviamente, multiplicar dos números largos también se puede reducir a multiplicar polinomios, así que dos números largos también se pueden multiplicar en tiempo $O(n \log n)$ (donde $n$ es la cantidad de dígitos de los números).

El descubrimiento de la transformada rápida de Fourier (FFT) se atribuye a Cooley y Tukey, que publicaron un algoritmo en 1965. Pero de hecho la FFT se había descubierto repetidamente antes, aunque su importancia no se entendió hasta la invención de las computadoras modernas. Algunos investigadores atribuyen el descubrimiento de la FFT a Runge y König en 1924. Pero en realidad Gauss ya desarrolló un método de este tipo en 1805, aunque nunca lo publicó.

Nótese que el algoritmo de FFT que se presenta aquí corre en tiempo $O(n \log n)$ , pero no sirve para multiplicar polinomios arbitrariamente grandes con coeficientes arbitrariamente grandes, ni para multiplicar enteros arbitrariamente grandes. Puede manejar sin problema polinomios de tamaño $10^5$ con coeficientes pequeños, o multiplicar dos números de tamaño $10^6$ , lo cual suele ser suficiente para resolver problemas de programación competitiva. Más allá de la escala de multiplicar números de $10^6$ bits, el rango y la precisión de los números de punto flotante usados durante el cálculo no serán suficientes para dar resultados finales exactos, aunque existen variantes más complejas que pueden realizar multiplicaciones de polinomios/enteros arbitrariamente grandes. Por ejemplo, en 1971 Schönhage y Strasser desarrollaron una variante para multiplicar números arbitrariamente grandes que aplica la FFT de forma recursiva en estructuras de anillos y corre en $O(n \log n \log \log n)$ . Y recientemente (en 2019) Harvey y van der Hoeven publicaron un algoritmo que corre en verdadero $O(n \log n)$ .

Transformada discreta de Fourier

Sea un polinomio de grado $n - 1$ :

$A(x) = a_0 x^0 + a_1 x^1 + \dots + a_{n-1} x^{n-1}$

Sin pérdida de generalidad asumimos que $n$ — la cantidad de coeficientes — es una potencia de $2$ . Si $n$ no es una potencia de $2$ , simplemente agregamos los términos faltantes $a_i x^i$ y ponemos los coeficientes $a_i$ en $0$ .

La teoría de los números complejos nos dice que la ecuación $x^n = 1$ tiene $n$ soluciones complejas (llamadas las raíces $n$ -ésimas de la unidad), y las soluciones son de la forma $w_{n, k} = e^{\frac{2 k \pi i}{n}}$ con $k = 0 \dots n-1$ . Además, estos números complejos tienen algunas propiedades muy interesantes: por ejemplo, la raíz $n$ -ésima principal $w_n = w_{n, 1} = e^{\frac{2 \pi i}{n}}$ se puede usar para describir todas las demás raíces $n$ -ésimas: $w_{n, k} = (w_n)^k$ .

La transformada discreta de Fourier (DFT) del polinomio $A(x)$ (o, equivalentemente, del vector de coeficientes $(a_0, a_1, \dots, a_{n-1})$ ) se define como los valores del polinomio en los puntos $x = w_{n, k}$ , es decir, es el vector:

$\begin{aligned} DFT (a_{0}, a_{1}, \dots, a_{n - 1}) & = (y_{0}, y_{1}, \dots, y_{n - 1}) \\ = (A (w_{n, 0}), A (w_{n, 1}), \dots, A (w_{n, n - 1})) \\ = (A (w_{n}^{0}), A (w_{n}^{1}), \dots, A (w_{n}^{n - 1})) \end{aligned}$

De forma similar se define la transformada discreta de Fourier inversa: la DFT inversa de los valores del polinomio $(y_0, y_1, \dots, y_{n-1})$ son los coeficientes del polinomio $(a_0, a_1, \dots, a_{n-1})$ .

$\text{InverseDFT}(y_0, y_1, \dots, y_{n-1}) = (a_0, a_1, \dots, a_{n-1})$

Así, si una DFT directa calcula los valores del polinomio en los puntos de las raíces $n$ -ésimas, la DFT inversa puede restaurar los coeficientes del polinomio usando esos valores.

Aplicación de la DFT: multiplicación rápida de polinomios

Sean dos polinomios $A$ y $B$ . Calculamos la DFT de cada uno: $\text{DFT}(A)$ y $\text{DFT}(B)$ .

¿Qué ocurre si multiplicamos estos polinomios? Obviamente, en cada punto los valores simplemente se multiplican, es decir

$(A \cdot B)(x) = A(x) \cdot B(x).$

Esto significa que si multiplicamos los vectores $\text{DFT}(A)$ y $\text{DFT}(B)$ — multiplicando cada elemento de un vector por el elemento correspondiente del otro — entonces obtenemos nada menos que la DFT del polinomio $\text{DFT}(A \cdot B)$ :

$\text{DFT}(A \cdot B) = \text{DFT}(A) \cdot \text{DFT}(B)$

Finalmente, aplicando la DFT inversa, obtenemos:

$A \cdot B = \text{InverseDFT}(\text{DFT}(A) \cdot \text{DFT}(B))$

En el lado derecho, el producto de las dos DFT se entiende como el producto componente a componente de los elementos de los vectores. Esto se puede calcular en tiempo $O(n)$ . Si podemos calcular la DFT y la DFT inversa en $O(n \log n)$ , entonces podemos calcular el producto de los dos polinomios (y, en consecuencia, también de dos números largos) con la misma complejidad temporal.

Hay que notar que los dos polinomios deben tener el mismo grado. De lo contrario, los dos vectores resultado de la DFT tienen distinta longitud. Podemos lograrlo agregando coeficientes con valor $0$ .

Y además, como el resultado del producto de dos polinomios es un polinomio de grado $2 (n - 1)$ , hay que duplicar los grados de cada polinomio (otra vez rellenando con $0$ s). A partir de un vector con $n$ valores no se puede reconstruir el polinomio deseado con $2n - 1$ coeficientes.

Transformada rápida de Fourier

La transformada rápida de Fourier es un método que permite calcular la DFT en tiempo $O(n \log n)$ . La idea básica de la FFT es aplicar divide y vencerás. Dividimos el vector de coeficientes del polinomio en dos vectores, calculamos recursivamente la DFT de cada uno y combinamos los resultados para calcular la DFT del polinomio completo.

Así, sea un polinomio $A(x)$ de grado $n - 1$ , donde $n$ es una potencia de $2$ y $n > 1$ :

$A(x) = a_0 x^0 + a_1 x^1 + \dots + a_{n-1} x^{n-1}$

Lo dividimos en dos polinomios más pequeños, uno que contiene solo los coeficientes de las posiciones pares y otro que contiene los coeficientes de las posiciones impares:

$\begin{aligned} A_{0} (x) & = a_{0} x^{0} + a_{2} x^{1} + \dots + a_{n - 2} x^{\frac{n}{2} - 1} \\ A_{1} (x) & = a_{1} x^{0} + a_{3} x^{1} + \dots + a_{n - 1} x^{\frac{n}{2} - 1} \end{aligned}$

Es fácil ver que

$A(x) = A_0(x^2) + x A_1(x^2).$

Los polinomios $A_0$ y $A_1$ tienen solo la mitad de coeficientes que el polinomio $A$ . Si podemos calcular $\text{DFT}(A)$ en tiempo lineal usando $\text{DFT}(A_0)$ y $\text{DFT}(A_1)$ , entonces obtenemos la recurrencia $T_{\text{DFT}}(n) = 2 T_{\text{DFT}}\left(\frac{n}{2}\right) + O(n)$ para la complejidad temporal, que resulta en $T_{\text{DFT}}(n) = O(n \log n)$ por el teorema maestro.

Veamos cómo se puede lograr eso.

Supongamos que ya calculamos los vectores $\left(y_k^0\right)$ y $\left(y_k^1\right)$ {k=0}^{n/2-1} = \text{DFT}(A_1) $(y_{k}^{1})_{k = 0}^{n /2 - 1} = DFT (A_{1})$ . Busquemos una expresión para $\left(y_k\right)_{k=0}^{n-1} = \text{DFT}(A)$ .

Para los primeros $\frac{n}{2}$ valores podemos usar simplemente la ecuación ya mencionada $A(x) = A_0(x^2) + x A_1(x^2)$ :

$y_k = y_k^0 + w_n^k y_k^1, \quad k = 0 \dots \frac{n}{2} - 1.$

Sin embargo, para los segundos $\frac{n}{2}$ valores necesitamos encontrar una expresión un poco distinta:

$\begin{aligned} y_{k + n / 2} & = A (w_{n}^{k + n / 2}) \\ = A_{0} (w_{n}^{2 k + n}) + w_{n}^{k + n / 2} A_{1} (w_{n}^{2 k + n}) \\ = A_{0} (w_{n}^{2 k} w_{n}^{n}) + w_{n}^{k} w_{n}^{n / 2} A_{1} (w_{n}^{2 k} w_{n}^{n}) \\ = A_{0} (w_{n}^{2 k}) - w_{n}^{k} A_{1} (w_{n}^{2 k}) \\ = y_{k}^{0} - w_{n}^{k} y_{k}^{1} \end{aligned}$

Aquí usamos de nuevo $A(x) = A_0(x^2) + x A_1(x^2)$ y las dos identidades $w_n^n = 1$ y $w_n^{n/2} = -1$ .

Por lo tanto obtenemos las fórmulas deseadas para calcular el vector completo $(y_k)$ :

$\begin{aligned} y_{k} & = y_{k}^{0} + w_{n}^{k} y_{k}^{1}, & k = 0 \dots \frac{n}{2} - 1, \\ y_{k + n / 2} & = y_{k}^{0} - w_{n}^{k} y_{k}^{1}, & k = 0 \dots \frac{n}{2} - 1. \end{aligned}$

(Este patrón $a + b$ y $a - b$ a veces se llama mariposa o butterfly.)

Así aprendimos a calcular la DFT en tiempo $O(n \log n)$ .

FFT inversa

Sea el vector $(y_0, y_1, \dots y_{n-1})$ — los valores del polinomio $A$ de grado $n - 1$ en los puntos $x = w_n^k$ — dado. Queremos restaurar los coeficientes $(a_0, a_1, \dots, a_{n-1})$ del polinomio. Este problema conocido se llama interpolación, y hay algoritmos generales para resolverlo. Pero en este caso especial (ya que conocemos los valores del polinomio en las raíces de la unidad) podemos obtener un algoritmo mucho más simple (que es prácticamente el mismo que la FFT directa).

Podemos escribir la DFT, según su definición, en forma matricial:

$(\begin{array}{cccccc} w_{n}^{0} & w_{n}^{0} & w_{n}^{0} & w_{n}^{0} & \dots & w_{n}^{0} \\ w_{n}^{0} & w_{n}^{1} & w_{n}^{2} & w_{n}^{3} & \dots & w_{n}^{n - 1} \\ w_{n}^{0} & w_{n}^{2} & w_{n}^{4} & w_{n}^{6} & \dots & w_{n}^{2 (n - 1)} \\ w_{n}^{0} & w_{n}^{3} & w_{n}^{6} & w_{n}^{9} & \dots & w_{n}^{3 (n - 1)} \\ ⋮ & ⋮ & ⋮ & ⋮ & ⋱ & ⋮ \\ w_{n}^{0} & w_{n}^{n - 1} & w_{n}^{2 (n - 1)} & w_{n}^{3 (n - 1)} & \dots & w_{n}^{(n - 1) (n - 1)} \end{array}) (\begin{array}{c} a_{0} \\ a_{1} \\ a_{2} \\ a_{3} \\ ⋮ \\ a_{n - 1} \end{array}) = (\begin{array}{c} y_{0} \\ y_{1} \\ y_{2} \\ y_{3} \\ ⋮ \\ y_{n - 1} \end{array})$

Esta matriz se llama matriz de Vandermonde.

Así, podemos calcular el vector $(a_0, a_1, \dots, a_{n-1})$ multiplicando el vector $(y_0, y_1, \dots y_{n-1})$ por la izquierda por la inversa de la matriz:

$(\begin{array}{c} a_{0} \\ a_{1} \\ a_{2} \\ a_{3} \\ ⋮ \\ a_{n - 1} \end{array}) = {(\begin{array}{cccccc} w_{n}^{0} & w_{n}^{0} & w_{n}^{0} & w_{n}^{0} & \dots & w_{n}^{0} \\ w_{n}^{0} & w_{n}^{1} & w_{n}^{2} & w_{n}^{3} & \dots & w_{n}^{n - 1} \\ w_{n}^{0} & w_{n}^{2} & w_{n}^{4} & w_{n}^{6} & \dots & w_{n}^{2 (n - 1)} \\ w_{n}^{0} & w_{n}^{3} & w_{n}^{6} & w_{n}^{9} & \dots & w_{n}^{3 (n - 1)} \\ ⋮ & ⋮ & ⋮ & ⋮ & ⋱ & ⋮ \\ w_{n}^{0} & w_{n}^{n - 1} & w_{n}^{2 (n - 1)} & w_{n}^{3 (n - 1)} & \dots & w_{n}^{(n - 1) (n - 1)} \end{array})}^{- 1} (\begin{array}{c} y_{0} \\ y_{1} \\ y_{2} \\ y_{3} \\ ⋮ \\ y_{n - 1} \end{array})$

Una comprobación rápida verifica que la inversa de la matriz tiene la siguiente forma:

$\frac{1}{n}$

Así obtenemos la fórmula:

$a_k = \frac{1}{n} \sum_{j=0}^{n-1} y_j w_n^{-k j}$

Comparando esto con la fórmula para $y_k$

$y_k = \sum_{j=0}^{n-1} a_j w_n^{k j},$

notamos que estos problemas son casi los mismos, así que los coeficientes $a_k$ se pueden encontrar con el mismo algoritmo de divide y vencerás, al igual que la FFT directa, solo que en lugar de $w_n^k$ hay que usar $w_n^{-k}$ , y al final hay que dividir los coeficientes resultantes por $n$ .

Así, el cálculo de la DFT inversa es casi el mismo que el de la DFT directa, y también se puede realizar en tiempo $O(n \log n)$ .

Implementación

Aquí presentamos una implementación recursiva simple de la FFT y de la FFT inversa, ambas en una sola función, porque la diferencia entre la FFT directa y la inversa es mínima. Para almacenar los números complejos usamos el tipo complex de la STL de C++.


using cd = complex<double>;
const double PI = acos(-1);

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

    vector<cd> 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);
    cd 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;
    }
}

A la función se le pasa un vector de coeficientes, y la función calculará la DFT o la DFT inversa y guardará el resultado otra vez en este vector. El argumento $\text{invert}$ indica si hay que calcular la DFT directa o la inversa. Dentro de la función primero comprobamos si la longitud del vector es igual a uno; si es el caso, no hay que hacer nada. En caso contrario, dividimos el vector $a$ en dos vectores $a0$ y $a1$ y calculamos la DFT de ambos de forma recursiva. Luego inicializamos el valor $wn$ y una variable $w$ , que contendrá la potencia actual de $wn$ . Después se calculan los valores de la DFT resultante usando las fórmulas de arriba.

Si el flag $\text{invert}$ está activado, reemplazamos $wn$ por $wn^{-1}$ , y cada uno de los valores del resultado se divide por $2$ (como esto se hace en cada nivel de la recursión, al final los valores quedarán divididos por $n$ ).

Usando esta función podemos crear una función para multiplicar dos polinomios:


vector<int> multiply(vector<int> const& a, vector<int> const& b) {
    vector<cd> fa(a.begin(), a.end()), fb(b.begin(), b.end());
    int n = 1;
    while (n < a.size() + b.size()) 
        n <<= 1;
    fa.resize(n);
    fb.resize(n);

    fft(fa, false);
    fft(fb, false);
    for (int i = 0; i < n; i++)
        fa[i] *= fb[i];
    fft(fa, true);

    vector<int> result(n);
    for (int i = 0; i < n; i++)
        result[i] = round(fa[i].real());
    return result;
}

Esta función trabaja con polinomios de coeficientes enteros, aunque también se puede adaptar para otros tipos. Como hay cierto error al trabajar con números complejos, hay que redondear los coeficientes resultantes al final.

Finalmente, la función para multiplicar dos números largos prácticamente no se diferencia de la función para multiplicar polinomios. Lo único que hay que hacer después es normalizar el número:


    int carry = 0;
    for (int i = 0; i < n; i++)
        result[i] += carry;
        carry = result[i] / 10;
        result[i] %= 10;
    }

Como la longitud del producto de dos números nunca supera la longitud total de ambos números, el tamaño del vector alcanza para realizar todas las operaciones de acarreo.

Implementación mejorada: cálculo in-place

Para aumentar la eficiencia pasaremos de la implementación recursiva a una iterativa. En la implementación recursiva de arriba separamos explícitamente el vector $a$ en dos vectores: los elementos en posiciones pares se asignaban a un vector temporal, y los de posiciones impares a otro. Sin embargo, si reordenamos los elementos de cierta forma, no hace falta crear esos vectores temporales (es decir, todos los cálculos se pueden hacer “in-place”, justo en el propio vector $A$ ).

Nótese que en el primer nivel de recursión, los elementos cuyo bit más bajo de la posición era cero se asignaron al vector $a_0$ , y los que tenían un uno como bit más bajo de la posición se asignaron a $a_1$ . En el segundo nivel de recursión ocurre lo mismo, pero con el segundo bit más bajo, etc. Por lo tanto, si invertimos los bits de la posición de cada coeficiente y los ordenamos por esos valores invertidos, obtenemos el orden deseado (se llama permutación de inversión de bits, bit-reversal permutation).

Por ejemplo, el orden deseado para $n = 8$ tiene la forma:

$a = \bigg{ \Big[ (a_0, a_4), (a_2, a_6) \Big], \Big[ (a_1, a_5), (a_3, a_7) \Big] \bigg}$

En efecto, en el primer nivel de recursión (rodeado por llaves), el vector se divide en dos partes $[a_0, a_2, a_4, a_6]$ y $[a_1, a_3, a_5, a_7]$ . Como se ve, en la permutación de inversión de bits esto corresponde simplemente a dividir el vector en dos mitades: los primeros $\frac{n}{2}$ elementos y los últimos $\frac{n}{2}$ elementos. Luego hay una llamada recursiva para cada mitad. Sean las DFT resultantes de cada una devueltas en el lugar de los propios elementos (es decir, la primera mitad y la segunda mitad del vector $a$ , respectivamente).

$a = \bigg{ \Big[y_0^0, y_1^0, y_2^0, y_3^0\Big], \Big[y_0^1, y_1^1, y_2^1, y_3^1 \Big] \bigg}$

Ahora queremos combinar las dos DFT en una para el vector completo. El orden de los elementos es ideal, y también podemos realizar la unión directamente en este vector. Podemos tomar los elementos $y_0^0$ e $y_0^1$ y aplicar la transformada mariposa. El lugar de los dos valores resultantes es el mismo que el de los dos valores iniciales, así que obtenemos:

$a = \bigg{ \Big[y_0^0 + w_n^0 y_0^1, y_1^0, y_2^0, y_3^0\Big], \Big[y_0^0 - w_n^0 y_0^1, y_1^1, y_2^1, y_3^1\Big] \bigg}$

De forma similar podemos calcular la transformada mariposa de $y_1^0$ e $y_1^1$ y poner los resultados en su lugar, y así sucesivamente. Como resultado obtenemos:

$a = \bigg{ \Big[y_0^0 + w_n^0 y_0^1, y_1^0 + w_n^1 y_1^1, y_2^0 + w_n^2 y_2^1, y_3^0 + w_n^3 y_3^1\Big], \Big[y_0^0 - w_n^0 y_0^1, y_1^0 - w_n^1 y_1^1, y_2^0 - w_n^2 y_2^1, y_3^0 - w_n^3 y_3^1\Big] \bigg}$

Así calculamos la DFT requerida a partir del vector $a$ .

Aquí describimos el proceso de calcular la DFT solo en el primer nivel de recursión, pero lo mismo funciona obviamente también para todos los demás niveles. Así, después de aplicar la permutación de inversión de bits, podemos calcular la DFT in-place, sin memoria adicional.

Esto además nos permite deshacernos de la recursión. Simplemente empezamos en el nivel más bajo, es decir, dividimos el vector en pares y les aplicamos la transformada mariposa. Esto deja el vector $a$ con el trabajo del último nivel aplicado. En el paso siguiente dividimos el vector en vectores de tamaño $4$ , y otra vez aplicamos la transformada mariposa, lo que nos da la DFT de cada bloque de tamaño $4$ . Y así sucesivamente. Finalmente, en el último paso obtenemos el resultado de las DFT de ambas mitades de $a$ , y aplicando la transformada mariposa obtenemos la DFT del vector $a$ completo.


using cd = complex<double>;
const double PI = acos(-1);

int reverse(int num, int lg_n) {
    int res = 0;
    for (int i = 0; i < lg_n; i++) {
        if (num & (1 << i))
            res |= 1 << (lg_n - 1 - i);
    }
    return res;
}

void fft(vector<cd> & a, bool invert) {
    int n = a.size();
    int lg_n = 0;
    while ((1 << lg_n) < n)
        lg_n++;

    for (int i = 0; i < n; i++) {
        if (i < reverse(i, lg_n))
            swap(a[i], a[reverse(i, lg_n)]);
    }

    for (int len = 2; len <= n; len <<= 1) {
        double ang = 2 * PI / len * (invert ? -1 : 1);
        cd wlen(cos(ang), sin(ang));
        for (int i = 0; i < n; i += len) {
            cd w(1);
            for (int j = 0; j < len / 2; j++) {
                cd u = a[i+j], v = a[i+j+len/2] * w;
                a[i+j] = u + v;
                a[i+j+len/2] = u - v;
                w *= wlen;
            }
        }
    }

    if (invert) {
        for (cd & x : a)
            x /= n;
    }
}

Primero aplicamos la permutación de inversión de bits intercambiando cada elemento con el elemento de la posición invertida. Luego, en los $\log n - 1$ estados del algoritmo, calculamos la DFT de cada bloque del tamaño correspondiente $\text{len}$ . Para todos esos bloques tenemos la misma raíz de la unidad $\text{wlen}$ . Iteramos todos los bloques y aplicamos la transformada mariposa en cada uno.

Podemos optimizar aún más la inversión de bits. En la implementación anterior iterábamos todos los bits del índice y creábamos el índice con los bits invertidos. Sin embargo, podemos invertir los bits de otra forma.

Supongamos que $j$ ya contiene el reverso de $i$ . Entonces, para pasar a $i + 1$ , hay que incrementar $i$ , y también hay que incrementar $j$ , pero en un sistema numérico “invertido”. Sumar uno en el sistema binario convencional equivale a invertir todos los unos de cola en ceros e invertir el cero justo antes de ellos en un uno. De forma equivalente, en el sistema numérico “invertido” invertimos todos los unos iniciales y también el siguiente cero.

Así obtenemos la siguiente implementación:


using cd = complex<double>;
const double PI = acos(-1);

void fft(vector<cd> & 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) {
        double ang = 2 * PI / len * (invert ? -1 : 1);
        cd wlen(cos(ang), sin(ang));
        for (int i = 0; i < n; i += len) {
            cd w(1);
            for (int j = 0; j < len / 2; j++) {
                cd u = a[i+j], v = a[i+j+len/2] * w;
                a[i+j] = u + v;
                a[i+j+len/2] = u - v;
                w *= wlen;
            }
        }
    }

    if (invert) {
        for (cd & x : a)
            x /= n;
    }
}

Además, podemos precomputar de antemano la permutación de inversión de bits. Esto es especialmente útil cuando el tamaño $n$ es el mismo para todas las llamadas. Pero incluso cuando solo tenemos tres llamadas (que son las necesarias para multiplicar dos polinomios), el efecto se nota. También podemos precomputar todas las raíces de la unidad y sus potencias.

Transformada teórica de números

Ahora cambiamos un poco el objetivo. Seguimos queriendo multiplicar dos polinomios en tiempo $O(n \log n)$ , pero esta vez queremos calcular los coeficientes módulo algún número primo $p$ . Por supuesto, para esta tarea podemos usar la DFT normal y aplicar el operador módulo al resultado. Sin embargo, hacer eso puede llevar a errores de redondeo, sobre todo al tratar con números grandes. La transformada teórica de números (NTT) tiene la ventaja de que solo trabaja con enteros y, por lo tanto, se garantiza que los resultados son correctos.

La transformada discreta de Fourier se basa en números complejos y en las raíces $n$ -ésimas de la unidad. Para calcularla de forma eficiente usamos de manera extensa propiedades de las raíces (por ejemplo, que hay una raíz que genera todas las demás por exponenciación).

Pero las mismas propiedades valen para las raíces $n$ -ésimas de la unidad en aritmética modular. Una raíz $n$ -ésima de la unidad en un cuerpo primo es un número $w_n$ que cumple:

$\begin{aligned} (w_{n})^{n} & = 1 (m o d p), \\ (w_{n})^{k} & \neq 1 (m o d p), 1 \leq k < n . \end{aligned}$

Las otras $n-1$ raíces se pueden obtener como potencias de la raíz $w_n$ .

Para aplicarlo en el algoritmo de la transformada rápida de Fourier, necesitamos que exista una raíz para algún $n$ que sea potencia de $2$ , y también para todas las potencias menores. Podemos notar la siguiente propiedad interesante:

$\begin{aligned} (w_{n}^{2})^{m} = w_{n}^{n} & = 1 (m o d p), with m = \frac{n}{2} \\ (w_{n}^{2})^{k} = w_{n}^{2 k} & \neq 1 (m o d p), 1 \leq k < m . \end{aligned}$

Así, si $w_n$ es una raíz $n$ -ésima de la unidad, entonces $w_n^2$ es una raíz $\frac{n}{2}$ -ésima de la unidad. Y en consecuencia, para todas las potencias de dos menores existen raíces del grado requerido, y se pueden calcular usando $w_n$ .

Para calcular la DFT inversa necesitamos el inverso $w_n^{-1}$ de $w_n$ . Pero para un módulo primo el inverso siempre existe.

Así, todas las propiedades que necesitamos de las raíces complejas también están disponibles en aritmética modular, siempre que tengamos un módulo $p$ lo bastante grande para el cual exista una raíz $n$ -ésima de la unidad.

Por ejemplo, podemos tomar los siguientes valores: módulo $p = 7340033$ , $w_{2^{20}} = 5$ . Si este módulo no alcanza, hay que encontrar otro par. Podemos usar el hecho de que para módulos de la forma $p = c 2^k + 1$ (y $p$ primo) siempre existe la raíz $2^k$ -ésima de la unidad. Se puede demostrar que $g^c$ es una raíz $2^k$ -ésima de la unidad de ese tipo, donde $g$ es una raíz primitiva de $p$ .


const int mod = 7340033;
const int root = 5;
const int root_1 = 4404020;
const int root_pw = 1 << 20;

void fft(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 ? root_1 : root;
        for (int i = len; i < root_pw; i <<= 1)
            wlen = (int)(1LL * wlen * wlen % mod);

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

    if (invert) {
        int n_1 = inverse(n, mod);
        for (int & x : a)
            x = (int)(1LL * x * n_1 % mod);
    }
}

Aquí la función inverse calcula el inverso modular (véase Inverso multiplicativo modular). Las constantes mod, root, root_pw determinan el módulo y la raíz, y root_1 es el inverso de root módulo mod.

En la práctica esta implementación es más lenta que la que usa números complejos (debido a la enorme cantidad de operaciones de módulo), pero tiene algunas ventajas como menor uso de memoria y ausencia de errores de redondeo.

Multiplicación con módulo arbitrario

Aquí queremos lograr el mismo objetivo que en la sección anterior. Multiplicar dos polinomios $A(x)$ y $B(x)$ , y calcular los coeficientes módulo algún número $M$ . La transformada teórica de números solo funciona para ciertos números primos. ¿Qué pasa cuando el módulo no tiene la forma deseada?

Una opción sería realizar varias transformadas teóricas de números con distintos números primos de la forma $c 2^k + 1$ , y luego aplicar el Teorema Chino del Resto para calcular los coeficientes finales.

Otra opción es repartir los polinomios $A(x)$ y $B(x)$ en dos polinomios más pequeños cada uno

$\begin{aligned} A (x) & = A_{1} (x) + A_{2} (x) \cdot C \\ B (x) & = B_{1} (x) + B_{2} (x) \cdot C \end{aligned}$

con $C \approx \sqrt{M}$ .

Entonces el producto de $A(x)$ y $B(x)$ se puede representar como:

$A(x) \cdot B(x) = A_1(x) \cdot B_1(x) + \left(A_1(x) \cdot B_2(x) + A_2(x) \cdot B_1(x)\right)\cdot C + \left(A_2(x) \cdot B_2(x)\right)\cdot C^2$

Los polinomios $A_1(x)$ , $A_2(x)$ , $B_1(x)$ y $B_2(x)$ contienen solo coeficientes menores que $\sqrt{M}$ ; por lo tanto, los coeficientes de todos los productos que aparecen son menores que $M \cdot n$ , lo cual suele ser lo bastante pequeño para manejarlo con tipos típicos de punto flotante.

Este enfoque, por lo tanto, requiere calcular los productos de polinomios con coeficientes más pequeños (usando la FFT normal y la FFT inversa), y luego el producto original se puede restaurar usando suma y multiplicación modular en tiempo $O(n)$ .

Aplicaciones

La DFT se puede usar en una enorme variedad de otros problemas que, a primera vista, no tienen nada que ver con multiplicar polinomios.

Todas las sumas posibles

Se nos dan dos arreglos $a[]$ y $b[]$ . Hay que encontrar todas las sumas posibles $a[i] + b[j]$ y, para cada suma, contar cuántas veces aparece.

Por ejemplo, para $a = [1,~ 2,~ 3]$ y $b = [2,~ 4]$ obtenemos: la suma $3$ se puede obtener de $1$ forma, la suma $4$ también de $1$ forma, $5$ de $2$ , $6$ de $1$ , $7$ de $1$ .

Construimos para los arreglos $a$ y $b$ dos polinomios $A$ y $B$ . Los números del arreglo actuarán como los exponentes en el polinomio ( $a[i] \Rightarrow x^{a[i]}$ ); y los coeficientes de ese término serán cuántas veces aparece el número en el arreglo.

Luego, al multiplicar estos dos polinomios en tiempo $O(n \log n)$ , obtenemos un polinomio $C$ , donde los exponentes nos dirán qué sumas se pueden obtener, y los coeficientes nos dirán cuántas veces. Para demostrarlo en el ejemplo:

$(1 x^1 + 1 x^2 + 1 x^3) (1 x^2 + 1 x^4) = 1 x^3 + 1 x^4 + 2 x^5 + 1 x^6 + 1 x^7$

Todos los productos escalares posibles

Se nos dan dos arreglos $a[]$ y $b[]$ de longitud $n$ . Hay que calcular los productos de $a$ con cada desplazamiento cíclico de $b$ .

Generamos dos arreglos nuevos de tamaño $2n$ : Invertimos $a$ y le agregamos $n$ ceros. Y simplemente concatenamos $b$ consigo mismo. Cuando multiplicamos estos dos arreglos como polinomios, y miramos los coeficientes $c[n-1],~ c[n],~ \dots,~ c[2n-2]$ del producto $c$ , obtenemos:

$c[k] = \sum_{i+j=k} a[i] b[j]$

Y como todos los elementos $a[i] = 0$ para $i \ge n$ :

$c[k] = \sum_{i=0}^{n-1} a[i] b[k-i]$

Es fácil ver que esta suma es justamente el producto escalar del vector $a$ con el $(k - (n - 1))$ -ésimo desplazamiento cíclico a la izquierda de $b$ . Así, estos coeficientes son la respuesta al problema, y aún así pudimos obtenerla en tiempo $O(n \log n)$ . Nótese aquí que $c[2n-1]$ también nos da el $n$ -ésimo desplazamiento cíclico, pero ese es el mismo que el $0$ -ésimo desplazamiento cíclico, así que no hace falta considerarlo por separado en la respuesta.

Dos franjas

Se nos dan dos franjas booleanas (arreglos cíclicos de valores $0$ y $1$ ) $a$ y $b$ . Queremos encontrar todas las formas de pegar la primera franja a la segunda de modo que en ninguna posición tengamos un $1$ de la primera franja al lado de un $1$ de la segunda.

El problema en realidad no se diferencia mucho del anterior. Pegar dos franjas significa simplemente que aplicamos un desplazamiento cíclico al segundo arreglo, y podemos pegar las dos franjas si el producto escalar de los dos arreglos es $0$ .

Matching de strings

Se nos dan dos strings, un texto $T$ y un patrón $P$ , formados por letras minúsculas. Hay que calcular todas las ocurrencias del patrón en el texto.

Creamos un polinomio para cada string ( $T[i]$ y $P[I]$ son números entre $0$ y $25$ correspondientes a las $26$ letras del alfabeto):

$A(x) = a_0 x^0 + a_1 x^1 + \dots + a_{n-1} x^{n-1}, \quad n = |T|$

con

$a_i = \cos(\alpha_i) + i \sin(\alpha_i), \quad \alpha_i = \frac{2 \pi T[i]}{26}.$

$B(x) = b_0 x^0 + b_1 x^1 + \dots + b_{m-1} x^{m-1}, \quad m = |P|$

con

$b_i = \cos(\beta_i) - i \sin(\beta_i), \quad \beta_i = \frac{2 \pi P[m-i-1]}{26}.$

Nótese que con la expresión $P[m-i-1]$ se invierte explícitamente el patrón.

Los coeficientes $(m-1+i)$ -ésimos del producto de los dos polinomios $C(x) = A(x) \cdot B(x)$ nos dirán si el patrón aparece en el texto en la posición $i$ .

$c_{m-1+i} = \sum_{j = 0}^{m-1} a_{i+j} \cdot b_{m-1-j} = \sum_{j=0}^{m-1} \left(\cos(\alpha_{i+j}) + i \sin(\alpha_{i+j})\right) \cdot \left(\cos(\beta_j) - i \sin(\beta_j)\right)$

con $\alpha_{i+j} = \frac{2 \pi T[i+j]}{26}$ y $\beta_j = \frac{2 \pi P[j]}{26}$

Si hay un match, entonces $T[i+j] = P[j]$ , y por lo tanto $\alpha_{i+j} = \beta_j$ . Esto da (usando la identidad trigonométrica de Pitágoras):

$\begin{aligned} c_{m - 1 + i} & = \sum_{j = 0}^{m - 1} (\cos (α_{i + j}) + i \sin (α_{i + j})) \cdot (\cos (α_{i + j}) - i \sin (α_{i + j})) \\ = \sum_{j = 0}^{m - 1} \cos (α_{i + j})^{2} + \sin (α_{i + j})^{2} = \sum_{j = 0}^{m - 1} 1 = m \end{aligned}$

Si no hay un match, entonces al menos un carácter es distinto, lo que hace que uno de los productos $a_{i+1} \cdot b_{m-1-j}$ no sea igual a $1$ , y eso lleva a que el coeficiente $c_{m-1+i} \ne m$ .

Matching de strings con comodines

Esta es una extensión del problema anterior. Esta vez permitimos que el patrón contenga el carácter comodín $*$ , que puede coincidir con cualquier letra posible. Por ejemplo, el patrón $a*c$ aparece en el texto $abccaacc$ en exactamente tres posiciones, en el índice $0$ , el índice $4$ y el índice $5$ .

Creamos exactamente los mismos polinomios, excepto que ponemos $b_i = 0$ si $P[m-i-1] = *$ . Si $x$ es la cantidad de comodines en $P$ , entonces habrá un match de $P$ en $T$ en el índice $i$ si $c_{m-1+i} = m - x$ .