---
# System prepended metadata

title: FFT 2D
tags: [Notas de clase, Tutoriales]

---

# FFT 2D

## Mutliplicación de polinomios en dos variables

Digamos que tenemos dos polinomios de dos variables $P(x,y)=\sum_{i=0}^{n_1}\sum_{j=0}^{m_1}p_{i,j} \cdot x^iy^j$ y $Q(x,y)=\sum_{i=0}^{n_2}\sum_{j=0}^{m_2}q_{i,j} \cdot x^iy^j$, y que queremos multiplicarlos, donde $n_1$ y $n_2$ son el máximo exponente de $x$ en $P$ y $Q$ respectivamente (lo que llamaremos grado en $x$), y $m_1$ y $m_2$ son el máximo exponente de $y$ en $P$ y $Q$ (lo que llamaremos grado en $y$). Esta multiplicación se puede hacer de forma bruta en complejidad  cuártica $O(n_1 \cdot m_1 \cdot n_2 \cdot m_2)$, iterando sobre cada par de términos de $P$ y $Q$.

En estas notas se va a explorar una forma de realizar esta multiplicación de forma eficiente en complejidad $O(nm\cdot log(nm))$.

## Prerrequisito: FFT
Antes de aplicar FFT para multiplicar polinomios de dos variables, es importante entender la idea que hay detrás de cómo multiplicar polinomios usando FFT. Para ello se pueden revisar las siguientes notas: https://hackmd.io/@alaneos777/H1KuU5ftw

## Representación en matriz
Así como un polinomio de una variable se puede representar en un arreglo, un polinomio $P$ de dos variables se puede representar en una matriz de tamaño $(n+1) \times (m+1)$, donde $n$ es el máximo exponente de $x$ en $P$ y $m$ es el máximo exponente de $y$ en $Q$.

El polinomio $P(x,y)=\sum_{i=0}^{n}\sum_{j=0}^{m}p_{i,j} \cdot x^iy^j$ se representaría con la matriz:

$$
\begin{bmatrix}
p_{0,0}&p_{0,1}&\dots&p_{0,m}\\
p_{1,0}&p_{1,1}&\dots&p_{1,m}\\
\vdots&\vdots&\ddots&\vdots\\
p_{n,0}&p_{n,1}&\dots&p_{n,m}\\
\end{bmatrix}
$$

## Multiplicación con representación punto-valor
Cuando aprendimos FFT, vimos que para multiplicar de forma rápida dos polinomios $P$ (de grado $n_1$) y $Q$ (de grado $n_2$) es posible utilizar su representación punto-valor. Recordemos que estos eran los pasos:

1. Evaluar $P$ y $Q$ en $n_1 + n_2 + 1$ valores distintos, a los que llamaremos $x_0, x_1, \dots x_{n_1+n_2}$, para obtener su representación punto-valor $[P(x_0), P(x_1), \dots, P(x_{n_1+n_2})]$ y $[Q(x_0), Q(x_1), \dots, Q(x_{n_1+n_2})]$.
2. Multiplicar punto a punto la representación punto-valor de los polinomios para obtener la representación punto-valor del producto de $P \cdot Q$, es decir, obtenemos $[P(x_0) \cdot Q(x_0), P(x_1) \cdot Q(x_1), \dots, P(x_{n_1+n_2}) \cdot Q(x_{n_1+n_2})]$.
3. Interpolamos la representación punto-valor del producto para obtener los coeficientes del producto.

El paso $2$ se puede realizar en tiempo lineal, y los pasos $1$ y $3$ dependen del método que se utilice para evaluar en distintos puntos e interpolar. La forma más eficiente es utilizar FFT para hacer la evaluación en distintos valores en tiempo $O(n \cdot \log(n))$, y la FFT inversa para interpolar el polinomio resultante en tiempo $O(n \cdot \log(n))$. 

Recordar que la FFT evalua el polinomio en $2^k$ valores distintos $\omega^0, \omega^1, \dots, \omega^{2^k-1}$, donde $\omega$ es la $(2^k)$-ésima raíz de la unidad, y la FFT inversa interpola un polinomio evaluado en dichos valores. Notar que $n_1+n_2+1$ no es necesariamente una potencia de dos, en ese caso elegimos una $k$ tal que $2^k \geq n_1+n_2+1$. Seguiremos obteniendo el producto correctamente aunque evaluemos los polinomios en más de $n_1+n_2+1$ valores distintos.

#### Multiplicando polinomios de dos variables con su representación punto-valor
Resulta que para multiplicar polinomios de dos variables se puede utilizar un método similar al anterior. Si queremos múltiplicar dos polinomios $P$ (con grado en $x$ igual a $n_1$ y grado en $y$ igual a $m_1$) y $Q$ (con grado en $x$ igual a $n_2$ y grado en $y$ igual a $m_2$), también se puede utilizar su representación punto-valor para multiplicarlos. El algoritmo sería el siguiente:

1. Elegir $n_1+n_2+1$ valores distintos $x_0, x_1, \dots, x_{n_1+n_2}$ y $m_1+m_2+1$ valores distintos $y_0, y_1, \dots, y_{m_1+m_2}$. Evaluar $P$ y $Q$ en las $(n_1+n_2+1) \times (m_1+m_2+1)$ parejas distintas de valores de $x$ y $y$ para obtener su representación punto-valor.

$$
\begin{bmatrix}
P(x_0, y_0)&P(x_0, y_1)&\dots& P(x_0, y_{m_1+m_2})\\
P(x_1, y_0)&P(x_1, y_1)&\dots& P(x_1, y_{m_1+m_2})\\
\vdots&\vdots&\ddots&\vdots\\
P(x_{n_1+n_2}, y_0)&P(x_{n_1+n_2}, y_1)&\dots &P(x_{n_1+n_2}, y_{m_1+m_2})\\
\end{bmatrix} \text{ y }
\begin{bmatrix}
Q(x_0, y_0)&Q(x_0, y_1)&\dots& Q(x_0, y_{m_1+m_2})\\
Q(x_1, y_0)&Q(x_1, y_1)&\dots& Q(x_1, y_{m_1+m_2})\\
\vdots&\vdots&\ddots&\vdots\\
Q(x_{n_1+n_2}, y_0)&Q(x_{n_1+n_2}, y_1)&\dots &Q(x_{n_1+n_2}, y_{m_1+m_2})\\
\end{bmatrix}
$$

2. Multiplicar punto a punto la representación punto-valor de los dos polinomios para obtener la representación punto-valor del producto.

$$
\begin{bmatrix}
(P\cdot Q)(x_0, y_0)&(P\cdot Q)(x_0, y_1)&\dots& (P\cdot Q)(x_0, y_{m_1+m_2})\\
(P\cdot Q)(x_1, y_0)&(P\cdot Q)(x_1, y_1)&\dots& (P\cdot Q)(x_1, y_{m_1+m_2})\\
\vdots&\vdots&\ddots&\vdots\\
(P\cdot Q)(x_{n_1+n_2}, y_0)&(P\cdot Q)(x_{n_1+n_2}, y_1)&\dots &(P\cdot Q)(x_{n_1+n_2}, y_{m_1+m_2})\\
\end{bmatrix}
$$

3. Interpolar la representación punto-valor del producto para obtener los coeficientes del polinomio resultante.

El paso $2$ se puede hacer en tiempo cuadrático. Ahora lo que nos falta ver es cómo realizar los pasos $1$ y $3$.

## FFT en un polinomio de dos variables
Podemos extender la idea de FFT en una variable a dos variables, en donde ahora tendremos $2^k$ valores distintos $\omega^0, \omega^1, \dots, \omega^{2^k-1}$ para evaluar $x$, y $2^l$ valores distintos $\psi^0, \psi^1, \dots, \psi^{2^l-1}$ para evaluar $y$, donde $\omega$ es la $(2^k)$-ésima raíz de la unidad y $\psi$ es la $(2^l)$-ésima ráiz de la unidad. Como el grado no siempre es potencia de $2$, elegimos la menor $k$ tal que $2^k \geq n_1+n_2+1$, y la menor $l$ tal que $2^l \geq m_1+m_2+1$.

La FFT 2D evaluaría $P(x,y)=\sum_{i=0}^{2^k-1}\sum_{j=0}^{2^l-1}p_{i,j} \cdot x^iy^j$ en todas las parejas de valores de $x$ y $y$:

$$
\begin{bmatrix}
p_{0,0}&p_{0,1}&\dots&p_{0,2^l-1}\\
p_{1,0}&p_{1,1}&\dots&p_{1,2^l-1}\\
\vdots&\vdots&\ddots&\vdots\\
p_{2^k-1,0}&p_{2^k-1,1}&\dots&p_{2^k-1,2^l-1}\\
\end{bmatrix}
\xrightarrow{\text{FFT 2D}}
\begin{bmatrix}
P(\omega^0, \psi^0)&P(\omega^0, \psi^1)&\dots& P(\omega^0, \psi^{2^l-1})\\
P(\omega^1, \psi^0)&P(\omega^1, \psi^1)&\dots& P(\omega^1, \psi^{2^l-1})\\
\vdots&\vdots&\ddots&\vdots\\
P(\omega^{2^k-1}, \psi^0)&P(\omega^{2^k-1}, \psi^1)&\dots& P(\omega^{2^k-1}, \psi^{2^l-1})\\
\end{bmatrix}
$$

:::warning
**Nota**
Si $P$ tiene grado en $x$ menor a $2^k-1$ o grado en $y$ menor a $2^l-1$, simplemente rellenamos los coeficientes restantes con ceros.
:::

Vamos a utilizar la FFT de una variable para implementar la FFT 2D. Para que sea esto posible, agrupemos los términos de $P$ de la siguiente forma:

$$
\begin{align}
P(x,y) &= \sum_{i=0}^{2^k-1}\sum_{j=0}^{2^l-1}p_{i,j} \cdot x^iy^j \\
    &= \sum_{i=0}^{2^k-1}x^i \sum_{j=0}^{2^l-1}p_{i,j} \cdot y^j \text{ (notemos que la segunda suma es un polinomio en $y$)}\\
    &= \sum_{i=0}^{2^k-1}x^i \cdot F_i(y)
\end{align}
$$

Donde $F_i$ es un polinomio de grado $2^l-1$:

$$
F_i(y) = \sum_{j=0}^{2^l-1}p_{i,j} \cdot y^j
$$

Si queremos evaluar a $P(x,y)$ en $(\omega^a,\psi^b)$, tenemos que:

$$
\begin{align}
P(\omega^a,\psi^b) &= \sum_{i=0}^{2^k-1}(\omega^a)^i \cdot F_i(\psi^b)
\end{align}
$$

Necesitamos evaluar a $F_i(y)$ para distintos valores $\psi^b$ ($0 \leq b < 2^l$), para lo cual ya tenemos un algoritmo: la FFT. Notemos que en la $i$-ésima fila de la representación en matriz de $P$ tenemos los coeficientes de $F_i$, por lo tanto, podemos hacer FFT en cada fila de dicha matriz para obtener todos los valores de $F_i(\psi^b)$:

$$
\begin{bmatrix}
p_{0,0}&p_{0,1}&\dots&p_{0,2^l-1}\\
p_{1,0}&p_{1,1}&\dots&p_{1,2^l-1}\\
\vdots&\vdots&\ddots&\vdots\\
p_{2^k-1,0}&p_{2^k-1,1}&\dots&p_{2^k-1,2^l-1}\\
\end{bmatrix}
\xrightarrow{\text{FFT por filas}}
\begin{bmatrix}
F_0(\psi^0)&F_0(\psi^1)&\dots& F_0(\psi^{2^l-1})\\
F_1(\psi^0)&F_1(\psi^1)&\dots& F_1(\psi^{2^l-1})\\
\vdots&\vdots&\ddots&\vdots\\
F_{2^k-1}(\psi^0)&F_{2^k-1}(\psi^1)&\dots& F_{2^k-1}(\psi^{2^l-1})\\
\end{bmatrix}
$$

Ya teniendo los valores de $F_i(\psi^b)$, podríamos obtener $P(\omega^a,\psi^b)$ en complejidad cúbica (para cada par de $(a,b)$, hacemos la suma de arriba en tiempo lineal). Pero podemos hacerlo todavía mejor. Ahora definamos el polinomio $C_b$ de la siguiente forma:

$$
\begin{align}
C_b(x) &= \sum_{i=0}^{2^k-1} F_i(\psi^b) \cdot x^i
\end{align}
$$

Se puede ver que $P(\omega^a,\psi^b)=C_b(\omega^a)$, y que los $F_i(\psi^b)$ ($0 \leq i < 2^k$) son los coeficientes de $C_b$. Ahora tenemos que evaluar $C_b(x)$ en distintos valores $\omega^0, \omega^1, \dots, \omega^{2^k-1}$. Después de aplicar la FFT por filas, ahora los coeficientes de $C_b$ están en la columna $b$, por lo tanto, si aplicamos FFT por columnas, obtendremos los valores de $C_b(\omega^a)$, que es lo mismo que obtener $P(\omega^a,\psi^b)$:

$$
\begin{bmatrix}
F_0(\psi^0)&F_0(\psi^1)&\dots& F_0(\psi^{2^l-1})\\
F_1(\psi^0)&F_1(\psi^1)&\dots& F_1(\psi^{2^l-1})\\
\vdots&\vdots&\ddots&\vdots\\
F_{2^k-1}(\psi^0)&F_{2^k-1}(\psi^1)&\dots& F_{2^k-1}(\psi^{2^l-1})\\
\end{bmatrix}
\xrightarrow{\text{FFT por columnas}}
\begin{bmatrix}
P(\omega^0, \psi^0)&P(\omega^0, \psi^1)&\dots& P(\omega^0, \psi^{2^l-1})\\
P(\omega^1, \psi^0)&P(\omega^1, \psi^1)&\dots& P(\omega^1, \psi^{2^l-1})\\
\vdots&\vdots&\ddots&\vdots\\
P(\omega^{2^k-1}, \psi^0)&P(\omega^{2^k-1}, \psi^1)&\dots& P(\omega^{2^k-1}, \psi^{2^l-1})\\
\end{bmatrix}
$$

Con esto está completa la FFT 2D sobre $P$. Primero hicimos $2^k$ FFTs sobre polinomios de grado $2^{l-1}$, y después $2^l$ FFTs sobre polinomios de grado $2^{k-1}$, por lo tanto la complejidad es $O(2^{k+l} \cdot \log (2^{k+l}))$, que se puede escribir también como $O(nm \cdot \log(nm))$ si $n$ es el grado en $x$ del producto y $m$ es el grado en $y$.

## FFT 2D inversa
Intuitivamente, podemos pensar que si para obtener la FFT 2D de un polinomio primero aplicamos FFT en las filas y luego en las columnas, entonces para hacer la FFT 2D inversa necesitamos aplicar la inversa de la FFT en las columnas y luego en las filas. Resulta que es cierto.

Si tenemos la representación punto-valor de un polinomio y queremos obtener sus coeficientes, primero aplicamos FFT inversa sobre las columnas para obtener los valores de $F_i(\psi^b)$:

$$
\begin{bmatrix}
P(\omega^0, \psi^0)&P(\omega^0, \psi^1)&\dots& P(\omega^0, \psi^{2^l-1})\\
P(\omega^1, \psi^0)&P(\omega^1, \psi^1)&\dots& P(\omega^1, \psi^{2^l-1})\\
\vdots&\vdots&\ddots&\vdots\\
P(\omega^{2^k-1}, \psi^0)&P(\omega^{2^k-1}, \psi^1)&\dots& P(\omega^{2^k-1}, \psi^{2^l-1})\\
\end{bmatrix}
\xrightarrow{\text{FFT inversa por columnas}}
\begin{bmatrix}
F_0(\psi^0)&F_0(\psi^1)&\dots& F_0(\psi^{2^l-1})\\
F_1(\psi^0)&F_1(\psi^1)&\dots& F_1(\psi^{2^l-1})\\
\vdots&\vdots&\ddots&\vdots\\
F_{2^k-1}(\psi^0)&F_{2^k-1}(\psi^1)&\dots& F_{2^k-1}(\psi^{2^l-1})\\
\end{bmatrix}
$$

Y después aplicamos FFT inversa por filas para finalmente obtener los coeficientes de $P$:

$$
\begin{bmatrix}
F_0(\psi^0)&F_0(\psi^1)&\dots& F_0(\psi^{2^l-1})\\
F_1(\psi^0)&F_1(\psi^1)&\dots& F_1(\psi^{2^l-1})\\
\vdots&\vdots&\ddots&\vdots\\
F_{2^k-1}(\psi^0)&F_{2^k-1}(\psi^1)&\dots& F_{2^k-1}(\psi^{2^l-1})\\
\end{bmatrix}
\xrightarrow{\text{FFT inversa por filas}}
\begin{bmatrix}
p_{0,0}&p_{0,1}&\dots&p_{0,2^l-1}\\
p_{1,0}&p_{1,1}&\dots&p_{1,2^l-1}\\
\vdots&\vdots&\ddots&\vdots\\
p_{2^k-1,0}&p_{2^k-1,1}&\dots&p_{2^k-1,2^l-1}\\
\end{bmatrix}
$$

La complejidad también será $O(nm \cdot \log(nm))$.

## Multiplicación de polinomios de más de dos variables.
Se pueden multiplicar polinomios de cualquier cantidad de variables utilizando un método similar al de dos variables. Necesitamos evaluar cada variable en potencias de la $(2^k)$-raíz de la unidad, donde $2^k$ es la menor potencia de $2$ que es mayor o igual al grado de dicha variable en el producto.

La FFT en $d$ dimensiones se realiza agrupando términos de igual forma que en 2D, es decir, se va a aplicar FFT de una variable por cada dimensión. Lo mismo con la FFT inversa.

Se muestra la representación gráfica de FFT sobre un polinomio de tres variables:

![FFT3D](https://hackmd.io/_uploads/ry75Y5jvC.png)

La complejidad quedaría como $O(d|P|\log(|P|)$, donde $d$ es la cantidad de variables y $|P|$ es el "tamaño" del polinomio $P$, donde definimos el tamaño como el producto del grado en cada variable más $1$.

## Observaciones
### NTT en 2 o más dimensiones
El mismo método que vimos para obtener la FFT 2D y su inversa, se puede aplicar para la NTT 2D y multiplicar polinomios de dos variables con coeficientes en $\mathbb{Z} / p\mathbb{Z}$ (módulo un primo).

También se puede usar NTT en más de dos dimensiones con la misma idea.

### XOR convolution en 2 o más dimensiones
También funciona con un método similar a este, aunque sólo he visto un problema que lo requería (y este era el tercero fácil de ese contest):

https://codeforces.com/gym/102978/problem/G

### ¿Importa el orden en el que se hace la FFT sobre las dimensiones?
Por la propiedad asociativa, conmutativa y distributiva de la suma y producto, no importa en qué orden se hacen las dimensiones al aplicar FFT en más de una dimensión. Por ejemplo, en el caso de dos variables, primero se puede hacer FFT sobre las columnas y luego las filas, y al hacer FFT inversa de nuevo primero sobre las columnas y luego sobre las filas.
