Нахождение обратных ко всем остаткам за $O(p)$



Часто бывает так, что в задаче нужно делить по модулю много раз. Это можно делать обычным алгоритмом взятия обратного по модулю за $O(\log p)$ на запрос. Если мы сделаем $n$ запросов, то асимптотика будет $O(n \log p)$. Сейчас мы рассмотрим алгоритм, который изначально предпосчитает обратные ко всем остаткам за $O(p)$, и тогда на запросы мы будем отвечать за $O(1)$. Если $n$ порядка $p$ или больше, то этот вариант будет более эффективен.

Есть много разных алгоритмов, которые делают это. Здесь будут представлены два, пожалуй, самых простых: один очень простой в понимании и написании, а другой еще легче в написании, однако не настолько очевидный с точки зрения понимания и придумывания.

Метод обратных факториалов

Идея первого алгоритма заключается в том, что мы посчитаем все возможные факториалы и обратные факториалы, а любое обратное к какому-то остатку представим как отношение двух факториалов.

Теорема: Теорема Вильсона гласит, что если $p$ — простое число, то

$$ (p - 1)! \equiv -1 \pmod p $$

Доказательство:

Давайте заметим, что все остатки от $1$ до $p - 1$ разбиваются на пары вида $x$, $x^{-1}$. Произведение чисел в паре равно единице по модулю $p$. Есть один крайний случай: когда $x = x^{-1}$. Это происходит в том случае, если $x^2 \equiv 1 \pmod p$, то есть $x^2 - 1 = (x - 1) \cdot (x + 1)\ ⋮\ p$. Значит, $x \equiv \pm 1 \pmod p$. Тогда в итоге $(p - 1)!$ по модулю $p$ состоит из произведения нескольких единиц, а также одной $-1$. Так что $(p - 1)! \equiv -1 \pmod p$. Что и требовалось доказать.

Зная этот факт, мы можем сразу понять, что $\left((p - 1)!\right)^{-1} \equiv -1 \pmod p$, потому что обратное к $-1$ — это $-1$. Таким образом, обратное к $(p - 1)!$ мы уже посчитали.

Замечание: На самом деле не обязательно было пользоваться этой формулой. Можно было посчитать за $O(p)$ число $(p - 1)!$, а потом бинарным возведением в степень найти к нему обратное за $O(\log p)$. Итоговая асимптотика бы от этого не пострадала.

Мы уже нашли обратное к $(p - 1)!$. Как же найти обратное к $(p - 2)!$ теперь? Заметим следующий факт:

$$ \frac{1}{k!} = \frac{k + 1}{(k + 1)!} $$

Так что алгоритм нахождения всех обратных факториалов следующий: идем с конца, изначально устанавливаем, что обратное к $(p - 1)!$ — это $-1$, а затем пересчитываем по очереди обратное к $k!$ как обратное к $(k + 1)!$, умноженное на $k + 1$.

Теперь пусть мы посчитали все факториалы и все обратные факториалы за $O(p)$. Как найти обратные ко всем остаткам? С этим нам поможет следующая формула:

$$ \frac{1}{k} = \frac{(k - 1)!}{k!} $$

А отношение двух факториалов — это произведение первого факториала на обратное ко второму.

Представим код алгоритма:

vector<int> get_all_modular_inverses(int p) {
    vector<int> inverse_factorials(p);
    inverse_factorials[p - 1] = p - 1; // -1 mod p = p - 1
    for (int k = p - 2; k > 0; k--) {
        inverse_factorials[k] = 1LL * inverse_factorials[k + 1] * (k + 1) % p;
    }

    vector<int> inverses(p);
    int factorial = 1;
    for (int k = 1; k < p; k++) {
        inverses[k] = 1LL * factorial * inverse_factorials[k] % p;
        factorial = 1LL * factorial * k % p;
    }
    return inverses;
}

Весь алгоритм — это два прохода по числам от $1$ до $p - 1$, так что работает он за $O(p)$. Потребление памяти тоже $O(p)$, потому что нужно хранить массивы факториалов и обратных факториалов. От одного из них можно избавиться, если вычислять ответ на лету (в приведенном коде мы не хранили факториалы), однако не от обоих.

Упражнение: Придумайте модернизацию этого алгоритма, которая работает за $O(p)$, но при этом потребляет $O(\sqrt{p})$ памяти (считайте, что ответы — вектор inverses — вы можете просто выводить на экран, и вам не нужно их хранить).

Замечание: Заметим, что можно считать обратные факториалы, начиная не обязательно с $p - 1$. Если нам нужно найти обратные ко всем остаткам от $1$ до $n$, то можно за $O(n)$ посчитать $n! \bmod p$, найти к нему обратное за $O(\log p)$ и потом аналогично представленному выше способу насчитать все обратные факториалы от $1$ до $n$. Тогда подсчет обратных ко всем остаткам от $1$ до $n$ будет работать за $O(n + \log p)$.

Как вы можете видеть, алгоритм очень простой. Однако его редко получится где-то применить, потому что, во-первых, нахождение всех обратных по отдельности работает за $O(p \log p)$, что тяжело отсечь от $O(p)$ на неучебной задаче, а во-вторых, модуль чаще всего — это число порядка $10^9$, поэтому вы не имеете возможности посчитать обратные ко всем остаткам, и использование стандартного алгоритма за $O(n \log p)$ дает более эффективное решение.

Алгоритм одного цикла

Второй алгоритм пишется всего одним циклом. Однако чтобы его вспомнить, придется написать пару формул на бумажке.

Алгоритм основывается на одном простом факте:

Теорема: $$ \frac{1}{k} \equiv -\left\lfloor \frac{p}{k} \right\rfloor \cdot \frac{1}{p \bmod k} \pmod p $$

Доказательство: Давайте представим $p$ в виде $k \cdot x + y$, где $x = \left\lfloor \frac{p}{k} \right\rfloor$ и $y = p \bmod k$.

Необходимо проверить, что

$$k \cdot \left(-\left\lfloor \frac{p}{k} \right\rfloor \cdot \frac{1}{p \bmod k} \right) \equiv 1 \pmod p$$

$$ k \cdot (-x \cdot \frac{1}{y}) = - (k \cdot x) \cdot \frac{1}{y} = - \left((k \cdot x + y) - y\right) \cdot \frac{1}{y} = - (p - y) \cdot \frac{1}{y} \equiv y \cdot \frac{1}{y} \equiv 1 \pmod p $$

Что и требовалось доказать.

Таким образом, мы можем посчитать обратное к $k$, если уже посчитано обратное к $p \bmod k$. Заметим, что это число меньше, чем $k$, поэтому все обратные можно вычислять по порядку.

Реализация у этого алгоритма крайне проста:

vector<int> get_all_modular_inverses(int p) {
    vector<int> inverses(p);
    inverses[1] = 1;
    for (int k = 2; k < p; k++) {
        inverses[k] = -1LL * (p / k) * inverses[p % k] % p + p;
        // +p because this number is negative
    }
    return inverses;
}

Также преимуществом этого метода является то, что это просто один цикл for по возрастанию, поэтому можно считать обратные не ко всем остаткам, а к первым $n$ остаткам за $O(n)$ очень легко. Однако не очень ясно, для чего это может вам понадобиться.

При тестировании на $p$ порядка $10^8$ второй алгоритм работает примерно в два раза быстрее, чем первый.