Spoiler: большая часть информации отсюда не имеет практического значения в решении задач, самая простая реализация FFT решает все необходимые задачи.
Примитивная реализация FFT обычно выглядит так (взято с емакса):
И в ней нет ничего плохого, однако мы хотим ее оптимизировать. Есть ряд действий, показывающих, как это сделать, но мы обсудим только те, которых нет в других источниках.
Посчитать заранее $$$\varepsilon_i$$$ для всех $$$i =1,\dots,n$$$. По-настоящему предпосчитать.
Сократить количество вызовов функции $$${\rm fft}$$$.
Избавиться от аллокации $$${\rm fb}$$$.
Прекальк корней из единицы в общем случае не работает из-за сопряжения, но сопряженная и не понадобится для этой реализации. Будем подсчитывать указанной ниже формулой:
ftype* wn;
uint32_t* rev;
const ld pi = acosl(-1.0);
void precalcRoots(uint32_t n) {
wn = new ftype[n];
wn[0] = {0,0}, wn[1] = {1,0};
for (uint32_t i = 1; (1 << i) <= n; ++i) {
ld ang = 2 * pi / (1 << (i + 1));
for (uint32_t j = 0; j < (1<<i); ++j)
wn[(1 << i) + j] = ftype(cosl(ang*j), sinl(ang*j));
}
}
Почему такое расположение корней из единицы? Заметим, что на $$$k'$$$ м снизу шаге рекурсии на $$$\ell'$$$ й позиции используется $$$\varepsilon_{\ell,k} = e^{\pm i \frac{2\pi\ell}{2k}}$$$ (да, двойку можно сократить, но так нагляднее). Таким образом, когда нам будет нужен $$$\varepsilon_{\ell,k},$$$ мы возьмем $$${\rm wn}[k+l]$$$.
Почему нам не понадобится сопряженная?
Утверждение: $$${\rm fft}({\rm fft}((a_0,\dots,a_{n-1}))) = (na_0,na_{n-1},\dots,na_1).$$$
Доказательство: Рассмотрим $$${\rm fft}({\rm fft}((a_0,\dots,a_{n-1}))) = (c_0,\dots,c_{n-1}).$$$ $$$c_k = \sum\limits_{r=0}^{n-1}\sum_{\ell=0}^{n-1}a_r\varepsilon_{\ell(k+r)} = \sum_{r=0}^{n-1}a_r\sum\limits_{\ell=0}^{n-1}\varepsilon_{\ell(k+r)}.$$$ При $$$k + r \vdots n$$$, а таких случаев ровно $$$n$$$, имеем $$$a_r$$$, в остальных 0 как следствие геометрической прогрессии. Отсюда и причина "разворота" коэффициентов.
Из того, что написано выше, следует, что достаточно просто посчитать $$${\rm fft}$$$ 2 раза, разделить результат на $$$n$$$ и развернуть коэффициенты, когда мы будем вносить их в результат.
Как избавиться от аллокации $$${\rm fb}$$$? На самом деле, кажется довольно странным, что мы размещаем коэффициенты многочленов в вещественные части, а мнимые оставляем нулями. Давайте сделаем следующее:
fa = new ftype[size];
for (uint32_t i = 0; i < size; i++)
fa[i] = ftype(i < n ? a[i] : 0, i < m ? b[i] : 0);
fft(fa, size);
for (uint32_t i=0;i<size;++i) fa[i] = fa[i]*fa[i];
fft(fa, size);
Мы поместили коэффициенты $$$A$$$ в вещественную часть и $$$B$$$ в мнимую, возвели многочлен в квадрат. Наш план действий: посчитать преобразование Фурье еще раз, чтобы получить коэффициенты многочлена $$${\rm fa}^2,$$$ затем поместить мнимые части коэффициентов этого многочлена в массив результата. Почему это и будет являться многочленом произведения $$$A$$$ и $$$B$$$?
Утверждение: Пусть $$$C(x) = \sum\limits_{k=0}^n(a_k + i\cdot b_k)x^k,\ {\rm fft}\left({\rm fft}\left(C^2\right)\right) = (nc_0,nc_n,\dots,nc_1), $$$ тогда $$${\rm Im}(c_k) = 2\sum\limits_{\ell=0}^k a_{\ell}b_{k-\ell}$$$.
Доказательство:
Рассмотрим мнимую часть при коэффициенте $$$c_k$$$, имеем $$$2\sum\limits_{\ell=0}^k a_{\ell}b_{k-\ell}$$$ (заметим, что в четном случае $$$c_{2k}$$$ имеем ровно одно вхождение $$$a_kb_k$$$, потому что в произведении $$$k < \ell$$$ строго).
Итак, мы внесли в fa коэффициенты, выполнили преобразование Фурье, возвели точки в квадрат, и снова выполнили преобразование Фурье. Теперь используем утверждения 1 и 2, чтобы получить ответ. Код для внесения значений в результат выглядит следующим образом:
uint64_t* res = new uint64_t[need];
res[0] = fa[0].im/(size<<1) + 0.5;
for (uint32_t i = 1; i < need; i++)
res[i] = fa[size-i].im/(size<<1) + 0.5;
dealloc();
return res;
Заметим, что делим на $$$2 \cdot {\rm size}$$$, так как при возведении в квадрат появляется биномиальный коэффициент.
Избавимся от рекурсии.
void fft(ftype* a, uint32_t n) {
assert((n & (n-1)) == 0);
for (uint32_t i = 1; i < n; i++)
if (i > rev[i])
swap(a[i], a[rev[i]]);
for (uint32_t len = 1; len < n; len <<= 1)
for (uint32_t i = 0; i < n; i += len<<1)
for (uint32_t j = 0; j < len; j++)
z = a[i+j+len] * wn[len+j], a[i+j+len] = a[i+j] - z, a[i+j] = a[i+j] + z;
}
Выражение n & (n-1) равняется нулю только тогда, когда $$$n$$$ — степень двойки.
Это все? И да, и нет. На самом деле, на сегодняшний день существует огромное количество разных алгоритмов, которые реализуют быстрое преобразование Фурье, и существуют целые библиотеки вроде FFTW, нацеленные на оптимальность. Можно реализовать split-radix FFT, но это займет пару тысяч строк, и ни одна олимпиадная задача не потребует таких усилий. На самом деле, в самом начале статьи было написано, что почти любая задача "загоняется" обычным рекурсивным FFT с тысячей аллокаций, но оптимизации в статье были приведены исключительно ради учебных целей. Может, когда-нибудь это кому-то поможет.
Касаемо бенчмарков: я запустил алгоритм с начала статьи, конечную его версию, и получил разницу в 1186 миллисекунд. Тесты проводились на 1000 чисел рандомного размера $$$n \leq 4000$$$, ввод/вывод осуществлялся одинаковым образом, из флагов стоял только O3. Первый алгоритм отработал за 1430 мс; конечная, оптимизированная версия — за 244 мс. Разница от количества запусков варьировалась не более, чем на 2 мс.
Полный код:
