В этом посте я подробно объясню принцип работы алгоритма Быстрого преобразования Фурье. Будет необходимо знание базовых операций с комплексными числами. Здесь приложена краткая реализация комплексных чисел.

/**< Будем представлять комплексное число парой (действительная часть, мнимая часть) */
typedef pair <long double,long double> comp;
const long double pi=2*acos(0.0);

/**< Сложение двух комплексных чисел */
comp operator +(comp a,comp b){
    comp c={a.first+b.first,a.second+b.second};
    return c;
}

/**< Разность двух комплексных чисел */
comp operator -(comp a,comp b){
    comp c={a.first-b.first,a.second-b.second};
    return c;
}

/**< Произведение двух комплексных чисел */
comp operator *(comp a,comp b){
    comp c={a.first*b.first-a.second*b.second,a.first*b.second+a.second*b.first};
    return c;
}

/**< Деление комплексного числа на действительное число */
comp operator /(comp a,long double b){
    comp c={a.first/b,a.second/b};
    return c;
}

Задача

Даны два полинома: $$$A = a_{0} + a_{1} \cdot x + \dots + a_{n-1} \cdot x^{n-1}$$$ и $$$B = b_{0} + b_{1} \cdot x + \dots + b_{n-1} \cdot x^{n-1}$$$. Для удобства будем записывать их как $$$[a_{0}, a_{1}, \cdot a_{n-1}]$$$ и $$$[b_{0}, b_{1}, \cdot b_{n-1}]$$$. Необходимо найти полином $$$C = A \cdot B$$$. Его размер равен $$$2n$$$.

Тривиальное решение

Задачу можно решить за асимптотику $$$O(n^{2})$$$, напрямую посчитав $$$C = A \cdot B = [a_{0} \cdot b_{0}, a_{0} \cdot b_{1} + a_{1} \cdot b_{0}, \dots ] = [\sum\limits_{i \in [0,0]}(a_{i} \cdot b_{0-i}), \sum\limits_{i \in [0,1]}(a_{i} \cdot b_{1-i}), \dots, \sum\limits_{i \in [0,k]}(a_{i} \cdot b_{k-i}), \dots]$$$.

План решения

Решение с помощью быстрого преобразования Фурье будет состоять из трех шагов:

1. Вычислить $$$A(x_{0}), A(x_{1}), \dots A(x_{2n-1})$$$ и $$$B(x_{0}), B(x_{1}), \dots B(x_{2n-1})$$$.

2. Вычислить значение $$$C$$$ в точках: $$$C(x_{0}) = A(x_{0}) \cdot B(x_{0}), C(x_{1}) = A(x_{1}) \cdot B(x_{1}), \dots C(x_{2n-1}) = A(x_{2n-1}) \cdot B(x_{2n-1})$$$.

3. Интерполировать $$$C$$$ по известным $$$2n$$$ значениям.

Вычисление значения полинома в точке в общем случае решается за $$$O(n)$$$. Поэтому первый шаг требует $$$O(n^{2})$$$ действий. Второй шаг решается одним проходом по массивам за $$$O(n)$$$. Интерполяция полинома решается в общем случае за $$$O(n^{2})$$$ с помощью интерполяционной формулы Лагранжа. Итоговая асимптотика решения для произвольных $$$x_{0}, x_{1}, \dots x_{2n-1}$$$ — $$$O(n^{2})$$$. Количество действий может быть существенно уменьшено, если выбрать $$$x_{0}, x_{1}, \dots x_{2n-1}$$$ особым образом.

Выбор множества для $$$X$$$

Рассмотрим множество $$$S$$$ комплексных чисел, модуль которых равен $$$1$$$. На комплексной плоскости ГМТ множества $$$S$$$ — окружность радиуса $$$1$$$ с центром в начале координат. Будем обозначать элемент множества $$$S$$$, аргумент которого равен $$$\phi$$$, как $$$\overline{\phi}$$$. Данное множество обладает замечательным свойством, на которое мы и будем опираться при написании алгоритма:

Теорема: Для $$$n \in \mathbb{Z}$$$ верно $$$(\overline{\phi})^{n} = \overline{n\phi}$$$.

Доказательство: докажем с помощью математической индукции. Для $$$n=0$$$ утверждение очевидно. Пусть утверждение верно для $$$n=k$$$. Докажем, что оно верно для $$$n=k+1$$$ и для $$$n=k-1$$$. $$$(\overline{\phi})^{k+1} = (\overline{\phi})^{k} \cdot \overline{\phi} = \overline{k\phi} \cdot \overline{\phi}$$$. Известно, что при перемножении комплексных чисел их модули перемножаются, а аргументы складываются. Поэтому $$$\overline{k\phi} \cdot \overline{\phi} = \overline {(k+1)\phi}$$$. Аналогично $$$(\overline{\phi})^{k-1} = \overline{k\phi} / \overline{\phi}$$$. Известно, что при делении комплексных чисел их модули делятся, а аргументы вычитаются. Поэтому $$$\overline{k\phi} / \overline{\phi} = \overline {(k-1)\phi}$$$.

Основной смысл этой теоремы — тот факт, что мы можем заменить степени на суммы. Кроме того, если выбрать в качестве множества $$$X$$$ точек числа $$$\overline{\frac{0 \cdot 2\pi}{2n}}, \overline{\frac{1 \cdot 2\pi}{2n}}, \dots \overline{\frac{(2n-1) \cdot 2\pi}{2n}}$$$, то любое число из этого множества в произвольной целой степени будет в этом множестве (говорят, множество $$$X$$$ замкнуто относительно операции целой степени).

Быстрое преобразование Фурье

Задача: дан многочлен $$$A = [a_{0}, a_{1}, \dots a_{n-1}]$$$. Необходимо вычислить $$$A(\overline{\frac{0 \cdot 2\pi}{n}}), A(\overline{\frac{1 \cdot 2\pi}{n}}), \dots A(\overline{\frac{(n-1) \cdot 2\pi}{n}})$$$. (Важно: выше мы вычисляли значения функция в $$$2n$$$ точках, так как нам необходимо знать $$$2n$$$ значений полинома $$$C$$$. Сейчас же мы вычислим значения в $$$n$$$ точках, а перед вызовом этого алгоритма просто допишем к $$$A$$$ $$$n$$$ нулей.) Для удобства будем считать $$$n$$$ целой степенью числа $$$2$$$.

Решение:

Будем решать задачу рекурсивно. Для $$$n=1$$$ нам необходимо вычислить массив из одного элемента $$$A(\overline{0}) = A(1) = a_{0}$$$. То есть, нам нужно вернуть массив без изменений.

Пусть теперь $$$n \ge 2$$$. Заметим, что $$$A(x) = a_{0}x^{0} + a_{1}x^{1} + \dots + a_{n-1}x^{n-1} = (a_{0} + a_{2}x^{2} + a_{4}x^{4} + \dots + a_{n-2}x^{n-2}) + x(a_{1} + a_{3}x^{2} + a_{5}x^{4} + \dots + a_{n-1}x^{n-2})$$$. Пусть $$$A_{1} = [a_{0}, a_{2}, a_{4}, \dots a_{n-2}]$$$ и $$$A_{2} = [a_{1}, a_{3}, a_{5}, \dots a_{n-1}]$$$. Тогда $$$A(x) = A_{1}(x^{2}) + x A_{2}(x^2)$$$. Кажется, что члены $$$x^{2k}$$$ усложняют задачу, но согласно теореме выше это $$$2kx$$$! Поэтому $$$A(x) = A_{1}(2x) + x A_{2}(2x)$$$.

Предположим, у нас есть алгоритм, решающий задачу для $$$n=\frac{k}{2}$$$. Решим задачу для $$$n=k$$$. Построим массивы $$$A_{1}$$$ и $$$A_{2}$$$ по определению выше и вызовем БПФ для них. Согласно предположению, мы умеем это делать. Теперь необходимо простроить БПФ для $$$A$$$. Пусть в массивах $$$B_{1}$$$ и $$$B_{2}$$$ лежат БПФ от них $$$A_{1}$$$ и $$$A_{2}$$$. Вспомним, что это $$$B_{1} = [A_{1}(\overline{\frac{0 \cdot 2\pi}{\frac{k}{2}}}), A_{1}(\overline{\frac{1 \cdot 2\pi}{\frac{k}{2}}}), \dots A_{1}(\overline{\frac{(\frac{k}{2}-1) \cdot 2\pi}{\frac{k}{2}}})]$$$ и $$$B_{2} = [A_{2}(\overline{\frac{0 \cdot 2\pi}{\frac{k}{2}}}), A_{2}(\overline{\frac{1 \cdot 2\pi}{\frac{k}{2}}}), \dots A_{2}(\overline{\frac{(\frac{k}{2}-1) \cdot 2\pi}{\frac{k}{2}}})]$$$. То есть, это значения полиномов $$$A_{1}$$$ и $$$A_{2}$$$ в $$$\frac{k}{2}$$$ точках равномерно распределенных по окружности комплексной плоскости. Пусть $$$B$$$ — результат БПФ для $$$A$$$. Заполним этот массив: $$$B[i] = A(\overline{\frac{i \cdot 2\pi}{k}}) = A_{1}(\overline{\frac{2i \cdot 2\pi}{k}}) + \overline{\frac{i \cdot 2\pi}{k}} \cdot A_{2}(\overline{\frac{2i \cdot 2\pi}{k}}) = A_{1}(\overline{\frac{i \cdot 2\pi}{\frac{k}{2}}}) + \overline{\frac{i \cdot 2\pi}{k}} \cdot A_{2}(\overline{\frac{i \cdot 2\pi}{\frac{k}{2}}}) = B_{1}[i] + \overline{\frac{i \cdot 2\pi}{k}} \cdot B_{2}[i]$$$. Для $$$i \ge \frac{k}{2}$$$ мы таким образом не можем посчитать $$$B[i]$$$, так как $$$B_{1}[i]$$$ и $$$B_{2}[i]$$$ мы явно не считали. Но заметим, что $$$B_{1}[i]=B_{1}[i+\frac{qk}{2}]$$$ и $$$B_{2}[i]=B_{2}[i+\frac{qk}{2}] \; q \in \mathbb{Z}$$$, то есть, значения функций периодические с периодом $$$T=\frac{k}{2}$$$. Но тогда для $$$i \ge \frac{k}{2}$$$ мы можем посчитать значения так: $$$B[i] = B_{1}[i-\frac{k}{2}] + \overline{\frac{i \cdot 2\pi}{k}} \cdot B_{2}[i-\frac{k}{2}]$$$, или для $$$i < \frac{k}{2}$$$ имеем $$$B[i+\frac{k}{2}] = B_{1}[i] + \overline{\frac{(i - \frac{k}{2}) \cdot 2\pi}{k}} \cdot B_{2}[i] = B_{1}[i] - \overline{\frac{i \cdot 2\pi}{k}} \cdot B_{2}[i]$$$, так как если отобразить аргумент комплексного числа, оно отобразится.

Асимптотика алгоритма $$$O(n \log{n})$$$.

/**< Будем передавать массивы по ссылке */
void fft(vector <comp> &a){
    int n=a.size();
    if(n==1) return; //Для полинома нулевой степени БПФ - он сам
    vector <comp> a1(n/2),a2(n/2); //Кладем, чередуя, коэффициенты в a1 и a2
    for(int i=0;i<n/2;i++)
        a1[i]=a[i*2],a2[i]=a[i*2+1];
    fft(a1),fft(a2);//Теперь в a1 и a2 - их БПФ
    for(int i=0;i<n/2;i++){
        comp x={cos(2*pi*i/n),sin(2*pi*i/n)}; //Это комплексное число с единичным модулем и аргументом равным 2*pi*i/n
        a[i]=a1[i]+x*a2[i];
        a[i+n/2]=a1[i]-x*a2[i];
    }
}

Обратное быстрое преобразование Фурье

Второй шаг плана выполняется за $$$O(n)$$$, необходимо просто одним проходом по массиву посчитать $$$C[i]=A[i] \cdot B[i]$$$.

Выполним теперь быструю интерполяцию. Заметим для начала, что уже написанный алгоритм БПФ решает следующую задачу: по данному массиву $$$A$$$ он считает:

$$$[w^{0} \quad w^{0} \quad w^{0} \dots \quad w^{0n-1}] \quad [a_{0}] \quad [x_{0}]$$$
$$$[w^{0} \quad w^{1} \quad w^{2} \dots \quad w^{1n-1}] \quad [a_{1}] \quad [x_{1}]$$$
$$$[w^{0} \quad w^{2} \quad w^{4} \dots \quad w^{2n-1}] \cdot [a_{2}] = [x_{2}]$$$
$$$[w^{0} \quad w^{3} \quad w^{6} \dots \quad w^{3n-1}] \quad [a_{3}] \quad [x_{3}]$$$
$$$\dots$$$
$$$[w^{0} \quad w^{n} \quad w^{2n} \dots \quad w^{n \cdot n -1}] \quad [a_{n-1}] \quad [x_{n-1}]$$$

, где $$$w = \overline{\frac{2\pi}{n}}$$$. (Это можно увидеть, перемножив матрицы и воспользовавшись тем, что $$$\overline{\phi}^{n} = \overline{n\phi}$$$.) Перепишем равенство в виде $$$W \cdot A = X$$$. Обратите внимание, что $$$W$$$ не зависит от содержимого $$$A$$$, для фиксированного $$$n$$$ она всегда одинакова. Теперь нам нужно по известному $$$X$$$ найти $$$A$$$. Предположим, что $$$W$$$ имеет обратную. Умножим равенство с двух сторон слева на обратную: $$$W^{-1} \cdot W \cdot X = W^{-1} \cdot A$$$, откуда $$$W^{-1} \cdot A = X$$$.

Утверждение: $$$W^{-1} = \frac{1}{n} \cdot$$$

$$$[w^{-0} \quad w^{-0} \quad w^{-0} \dots \quad w^{-0n+1}]$$$
$$$[w^{-0} \quad w^{-1} \quad w^{-2} \dots \quad w^{-1n+1}]$$$
$$$[w^{-0} \quad w^{-2} \quad w^{-4} \dots \quad w^{-2n+1}]$$$
$$$[w^{-0} \quad w^{-3} \quad w^{-6} \dots \quad w^{-3n+1}]$$$
$$$\dots$$$
$$$[w^{-0} \quad w^{-n} \quad w^{-2n} \dots \quad w^{-n \cdot n +1}]$$$

Доказательство: Нам необходимо показать, что $$$W \cdot W^{-1} = I$$$. Перемножим строку $$$i$$$ матрицы $$$W$$$ со столбцом $$$j$$$ матрицы $$$nW^{-1}$$$, так, что $$$i \neq j$$$. Имеем: $$$w^{0}w^{0} + w^{i}w^{-j} + w^{2i}w^{-2j} + \dots + w^{(n-1)i}w^{(n-1)j} = \overline{\frac{0(i-j) \cdot 2\pi}{n}} + \overline{\frac{1(i-j) \cdot 2\pi}{n}} + \overline{\frac{2(i-j) \cdot 2\pi}{n}} + \dots + \overline{\frac{(n-1)(i-j) \cdot 2\pi}{n}}$$$. Но это четное количество комплексных чисел, размещенных на единичной окружности с равным шагом, поэтому их сумма равна $$$0$$$. Перемножим теперь строку $$$i$$$ матрицы $$$W$$$ со столбцом $$$i$$$ матрицы $$$nW^{-1}$$$. Имеем $$$w^{0}w^{0} + w^{i}w^{-i} + w^{2i}w^{-2i} + \dots + w^{(n-1)i}w^{(n-1)i} = \overline{\frac{0 \cdot 2\pi}{n}} + \overline{\frac{0 \cdot 2\pi}{n}} + \overline{\frac{0 \cdot 2\pi}{n}} + \dots + \overline{\frac{0 \cdot 2\pi}{n}} = n$$$.

Будем искать $$$nW^{-1}X = nA$$$, а затем поделим результат на $$$n$$$. Тот факт, что теперь все степени в матрице $$$nW^{-1}$$$ отображены означает, что и значение $$$x$$$ в формуле $$$B[i] = B_{1}[i] + \overline{\frac{i \cdot 2\pi}{k}} \cdot B_{2}[i]$$$ должен поменять знак. Так как $$$B_{1}$$$ и $$$B_{2}$$$ содержат в качестве аргумента $$$x^{2}$$$, их значение не поменяется. Тогда $$$B[i] = B_{1}[i] + \overline{-\frac{i \cdot 2\pi}{k}} \cdot B_{2}[i]$$$. Аналогично $$$B[i+\frac{n}{2}] = B_{1}[i] - \overline{-\frac{i \cdot 2\pi}{k}} \cdot B_{2}[i]$$$.

Асимптотика алгоритма $$$O(n \log{n})$$$.

void revfft(vector <comp> &a){
    int n=a.size();
    if(n==1) return;
    vector <comp> a1(n/2),a2(n/2);
    for(int i=0;i<n/2;i++)
        a1[i]=a[i*2],a2[i]=a[i*2+1];
    revfft(a1),revfft(a2);
    for(int i=0;i<n/2;i++){
        comp x={cos(2*pi*i/n),-sin(2*pi*i/n)}; //Это комплексное число с единичным модулем и аргументом равным -2*pi*i/n
        a[i]=a1[i]+x*a2[i];
        a[i+n/2]=a1[i]-x*a2[i];
        a[i]=a[i]/2; //Глубина рекурсии lb(n), поэтому мы поделим a на 2 всего lb(n) раз, в результате он уменьшится в n раз
        a[i+n/2]=a[i+n/2]/2;
    }
}

Умножение полиномов

Теперь мы готовы написать умножение полиномов за $$$O(n \log{n})$$$.

vector <int> mult(vector <int> a,vector <int> b){
    /**< Подготовка полиномов */
    vector <comp> fa,fb; //Переводим числа в комплексные
    for(int i=0;i<a.size();i++) fa.push_back({a[i],0});
    for(int i=0;i<b.size();i++) fb.push_back({b[i],0});
    int l=0;
    while(1<<l < a.size()+b.size())
        l++;
    int n=1<<l; //n - наименьшая целая степень числа 2 не меньшая, чем сумма длин данных полиномов
    while(fa.size()<n*2) fa.push_back({0,0}); // Дополняем массивы до размера 2n
    while(fb.size()<n*2) fb.push_back({0,0});

    fft(fa),fft(fb); //Шаг 1
    for(int i=0;i<n*2;i++) //Шаг 2
        fa[i]=fa[i]*fb[i];
    revfft(fa); //Шаг 3

    vector <int> res(fa.size()); // Переводим комплексные числа в целые
    for(int i=0;i<fa.size();i++)
        res[i]=(int)(fa[i].first+0.5);
    while(!res.empty() && res.back()==0) res.pop_back();
    return res;
}

Улучшение алгоритма БПФ

Полученный алгоритм сильно блуждает по памяти, из за чего необходимые значения плохо кэшируются. Можно алгоритм сильно ускорить, изменив принцип обхода. Посмотрим, как в процессе вычисления БПФ делится наш полином. На первом этапе в первый полином попадают числа, чей последний бит в индексе равен $$$0$$$, а во второй полином — чей бит равен $$$1$$$. На втором этапе числа делятся по второму с конца биту. Иными словами, если мы отсортируем числа в зеркально написанной двоичной форме, массив на каждом этапе будет делится ровно на две части: левая половина уйдет в первый массив, а правая — во второй.

Для $$$n=4$$$, эта перестановка: $$$[0,2,1,3]$$$, для $$$n=8$$$, эта перестановка: $$$[0,4,2,6,1,5,3,7]$$$.

Отсортировать числа по зеркально написанной двоичной форме можно за $$$O(n)$$$. После этого мы можем избавиться от рекурсии. Для начала разделим массив на отрезки длины $$$2$$$ и посчитаем БПФ от них, запишем результат в них же. Затем разделим массив на отрезки длины $$$4$$$ и повторим процесс. Когда мы на очередном этапе хотим посчитать БПФ на отрезке $$$[l,r)$$$, в отрезках $$$[l,m)$$$ и $$$[m,r)$$$ уже записаны их БПФ.

void fft(vector <comp> &a){
    int n=a.size();

    /**< Сортировка по двоичной форме, записанной зеркально */
    for(int i=1,j=0;i<n;i++){
        int bit=(n>>1);
        while(j>=bit){
            j-=bit;
            bit>>=1;
        }
        j+=bit;
        if(i<j)
            swap(a[i],a[j]);
    }

    for(int len=2;len<=n;len<<=1){ //Строим итеративно БПФ для отрезков длины 2, 4, 8, ... n
        comp w={cos(2*pi/len),sin(2*pi/len)}; //Множитель, умножив на который, мы прокрутим аргумент на один пункт длины 2*pi/len дальше
        for(int i=0;i<n;i+=len){ //Обрабатываем поблочно, i - начало очередного блока
            comp cur_w={1,0}; //Аргумент
            for(int j=0;j<len/2;j++){ //Итерируемся по блоку
                comp u=a[i+j],v=a[i+j+len/2]*cur_w; //Так как записываем в эти элементы массива, их нужно скопировать
                a[i+j]=u+v;
                a[i+j+len/2]=u-v;
                cur_w=cur_w*w; //Прокручиваем аргумент на угол 2*pi/len
            } //Когда мы вышли из цикла, аргумент cur_w снова равен нулю, так как мы добавили 2*pi/len len раз.
        }
    }
}

void revfft(vector <comp> &a){
    int n=a.size();

    /**< Сортировка по двоичной форме, записанной зеркально */
    for(int i=1,j=0;i<n;i++){
        int bit=(n>>1);
        while(j>=bit){
            j-=bit;
            bit>>=1;
        }
        j+=bit;
        if(i<j)
            swap(a[i],a[j]);
    }

    for(int len=2;len<=n;len<<=1){
        comp w={cos(2*pi/len),-sin(2*pi/len)}; //Значения аргумента отображены
        for(int i=0;i<n;i+=len){
            comp cur_w={1,0};
            for(int j=0;j<len/2;j++){
                comp u=a[i+j],v=a[i+j+len/2]*cur_w;
                a[i+j]=u+v;
                a[i+j+len/2]=u-v;
                cur_w=cur_w*w;
            }
        }
    }

    for(int i=0;i<n;i++) //Мы вычислили nW^{-1}, нужно разделить массив на n
        a[i]=a[i]/n;
}

Так как мы $$$\log{n}$$$ раз проходимся по всему массиву без прыжков, данные хорошо кэшируются, и алгоритм ускоряется в $$$\sim 10$$$ раз.

Задача