Эффективная реализация умножения Карацубы с авто-векторизацией

Правка ru2, от dmkz, 2020-02-24 12:07:38

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

$$$\text{ }$$$

Всем привет!

Давным-давно меня попросил один человек научить его находить миллионное число Фибоначчи абсолютно точно на 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) + a_1(x)) \cdot (b_0(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.

Стресс-тест на 512К для 64-битных и 32-битных коэффициентов

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

Теги карацуба, умножение, многочлены

История

 
 
 
 
Правки
 
 
  Rev. Язык Кто Когда Δ Комментарий
en2 Английский dmkz 2020-02-24 12:39:36 1 Tiny change: 'l position will be g' -> 'l positions will be g'
ru3 Русский dmkz 2020-02-24 12:38:29 79
en1 Английский dmkz 2020-02-24 12:36:53 5579 Initial revision for English translation
ru2 Русский dmkz 2020-02-24 12:07:38 9 Мелкая правка: '(a_0(x) + b_0(x)) \cdot (a_1(x) + b_1(' -> '(a_0(x) + a_1(x)) \cdot (b_0(x) + b_1('
ru1 Русский dmkz 2020-02-24 12:05:14 5530 Первая редакция (опубликовано)