Совсем недавно у меня возникла необходимость решить такую задачу:
Для заданных $$$a$$$ и $$$m$$$ найти все такие $$$x$$$, что: $$$x^2 \equiv a \pmod{m}$$$.
Спустя некоторое время поиска по интернету и тщательному разбору каждого случая, я получил решение, которое работает для любых $$$a$$$ и $$$m > 0$$$ за время $$$O(\sqrt{m})$$$, либо за время $$$O(\log^2(m))$$$ если $$$a$$$ и $$$m$$$ взаимнопростые.
Ниже приведу рассмотренные мной случаи и ссылки на статьи (или обсуждения) для алгоритмов:
- $$$x^2 \equiv a \pmod{p}$$$, $$$a \ne 0$$$, $$$p$$$ — нечетное простое (ссылка).
- $$$x^2 \equiv a \pmod{p^k}$$$, $$$a \ne 0$$$, $$$k > 0$$$, $$$\gcd(a, p) = 1$$$, $$$p$$$ — нечетное простое (ссылка).
- $$$x^2 \equiv a \pmod{2^k}$$$, $$$a \ne 0$$$, $$$k > 0$$$, $$$\gcd(a, 2) = 1$$$ (cсылка).
- $$$x^2 \equiv 0 \pmod{p^k}$$$, $$$k > 0$$$, $$$p$$$ — простое. Заметим, что в таком случае $$$x^2$$$ должен делиться на $$$p^k$$$. Тогда $$$x$$$ должен делиться на $$$p^{\lceil\frac{k}{2}\rceil}$$$. Отсюда получаем, что $$$x = i \cdot p^{\lceil\frac{k}{2}\rceil}$$$, $$$0 \le i < p^{\lfloor\frac{k}{2}\rfloor}$$$.
- $$$x^2 \equiv p^s \pmod{p^k}$$$, $$$k > 0$$$, $$$0 < s < k$$$, $$$p$$$ — простое. Решения существуют только для $$$s = 2 \cdot t$$$. Доказывается от противного. Тогда очевидно, что $$$p^t$$$ является одним из решений ($$$(p^t)^2 = p^s$$$). Найдем оставшиеся решения. Пусть $$$x = p^t + y$$$, тогда $$$x^2 = p^s + 2 \cdot p^t \cdot y + y^2$$$. Заметим, что $$$2 \cdot p^t \cdot y + y^2$$$ должен делиться на $$$p^k$$$. В таком случае $$$x = p^t + i \cdot p^{\max(\lceil\frac{k}{2}\rceil, k - t - (p = 2))}$$$.
- $$$x^2 \equiv q \cdot p^s \pmod{p^k}$$$, $$$k > 0$$$, $$$0 < s < k$$$, $$$\gcd(q, p) = 1$$$, $$$p$$$ — простое. Пусть $$$x_1$$$ решение для $$$x^2 \equiv q \pmod{p^k}$$$, а $$$x_2$$$ решение для $$$x^2 \equiv p^s \pmod{p^k}$$$. Тогда $$$x = x_1 \cdot x_2$$$ будет являться решением исходной задачи.
Последние три случая пришлось выводить самостоятельно, поэтому ссылок нет. Доказательства пока поленился оформить в этом посте, хотя в целом они не очень сложны.
Общее решение находится с использованием Китайской теоремы об остатках и сведению к одному из указанных выше случаев.
Также приведу полную реализацию на Python (функция, решающая общий случай — discrete_sqrt):
. Поэтому нам достаточно перебрать первые 
, по памяти: 
, по памяти: 
, но если предподсчитать 
, то слияние можно делать за 
. Заметим, что слияние можем делать сразу степенями двойки, но ограничения задачи не требовали этого.
