В данном блоге приводится эффективная реализация умножения Карацубы двух многочленов.

Всем привет!

Давным-давно меня попросил один человек научить его находить миллионное число Фибоначчи абсолютно точно на C++ за секунду. Эта задача была успешно решена умножением Карацубы. После этого я попробовал сдавать стандартные задачи вроде "сопоставить одну строку из символов $$$\text{ATGC}$$$ к каждой позиции другой строки и найти позицию с максимальным числом совпадений" алгоритмом Карацубы, и они успешно сдавались с двукратным запасом по времени. Было много реализаций, но в итоге MrDindows помог написать самую эффективную из всех, которые придумывались, и я решил поделиться ей.

Идея умножения Карацубы

Пусть у нас есть два многочлена $$$a(x)$$$ и $$$b(x)$$$ равной длины $$$2n$$$ и мы хотим их умножить. Представим их как $$$a(x) = a_0(x) + x^n \cdot a_1(x)$$$ и $$$b(x) = b_0(x) + x^n \cdot b_1(x)$$$. Теперь посчитаем результат умножения $$$a(x) \cdot b(x)$$$ следующим образом:

• Вычислим $$$E(x) = (a_0(x) + b_0(x)) \cdot (a_1(x) + b_1(x))$$$
• Вычислим $$$r_0(x)=a_0(x)\cdot b_0(x)$$$
• Вычислим $$$r_2(x)=a_1(x) \cdot b_1(x)$$$
• Тогда $$$a(x) \cdot b(x) = r_0(x) + x^n \cdot \left(E(x) - r_0(x) - r_2(x)\right) + x^{2n} \cdot r_2(x)$$$

Видим, что вместо четырех умножений, мы сделали три умножения многочленов вдвое меньшей длины, поэтому асимптотика получается $$$O(n^{\log_2(3)})$$$ или $$$O(n^{1.58})$$$

Эффективная реализация

#pragma GCC optimize("Ofast")
#pragma GCC target("avx,avx2,fma") 
namespace {
    template<int n, typename T>
    void mult(const T *__restrict a, const T *__restrict b, T *__restrict res) {
        if (n <= 64) { // если длина маленькая, то эффективнее работает наивное умножение за квадрат
            for (int i = 0; i < n; i++) {
                for (int j = 0; j < n; j++) {
                    res[i + j] += a[i] * b[j];
                }
            }
        } else {
            const int mid = n / 2;
            alignas(64) T btmp[n], E[n] = {};
            auto atmp = btmp + mid;
            for (int i = 0; i < mid; i++) {
                atmp[i] = a[i] + a[i + mid]; // atmp(x) - сумма двух половинок многочлена a(x)
                btmp[i] = b[i] + b[i + mid]; // btmp(x) - сумма двух половинок многочлена b(x)
            }
            mult<mid>(atmp, btmp, E); // вычисляем E(x) = (alow(x) + ahigh(x)) * (blow(x) + bhigh(x))
            mult<mid>(a + 0, b + 0, res); // вычисляем rlow(x) = alow(x) * blow(x)
            mult<mid>(a + mid, b + mid, res + n); // вычисляем rhigh(x) = ahigh(x) * bhigh(x)
            for (int i = 0; i < mid; i++) { // теперь считаем rmid(x) = E(x) - rlow(x) - rhigh(x) и сразу записываем
                const auto tmp = res[i + mid];
                res[i + mid] += E[i] - res[i] - res[i + 2 * mid];
                res[i + 2 * mid] += E[i + mid] - tmp - res[i + 3 * mid];
            }
        }
    }
}

Небольшое тестирование

Пример решения реальных задач: 528D. Нечёткий поиск и AtCoder Beginner Contest 149 E. Handshake.

#pragma GCC optimize("Ofast,unroll-loops")
#pragma GCC target("avx,avx2,fma")
#include <bits/stdc++.h>

namespace {
    template<int n, typename T>
    void mult(const T *__restrict a, const T *__restrict b, T *__restrict res) {
        if (n <= 64) {
            for (int i = 0; i < n; i++) {
                for (int j = 0; j < n; j++) {
                    res[i + j] += a[i] * b[j];
                }
            }
        } else {
            const int mid = n / 2;
            alignas(64) T btmp[n], E[n] = {};
            auto atmp = btmp + mid;
            for (int i = 0; i < mid; i++) {
                atmp[i] = a[i] + a[i + mid];
                btmp[i] = b[i] + b[i + mid];
            }
            mult<mid>(atmp, btmp, E);
            mult<mid>(a + 0, b + 0, res);
            mult<mid>(a + mid, b + mid, res + n);
            for (int i = 0; i < mid; i++) {
                const auto tmp = res[i + mid];
                res[i + mid] += E[i] - res[i] - res[i + 2 * mid];
                res[i + 2 * mid] += E[i + mid] - tmp - res[i + 3 * mid];
            }
        }
    }
}

using ll = long long;

int main() {
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(0);
    const int NMAX = (1 << 19);
    alignas(64) static ll a[NMAX], b[NMAX], c[NMAX * 2];
    for (int i = 0; i < NMAX; i++) a[i] = b[i] = 1;
    mult<NMAX>(a,b,c);
    ll sum = 0;
    for (int i = 0; i < 2 * NMAX; i++) sum += c[i];
    std::cout << sum << std::endl;
    return 0;
}

Возможно, это даже можно где-нибудь использовать, но будьте осторожны, так как для умножения двух многочленов длины $$$2^{19}$$$ алгоритм уже выполняет $$$3.047.478.784$$$ элементарных операций.