Квадратный корень по простому модулю за $O(\log p)$
Задача формулируется следующим образом. Даны числа $a$ и $p$. При этом $p$ — простое. Нужно найти такое $z$, что $z^2 \bmod p = a$ или сказать, что такого $z$ не существует.
Давайте опишем алгоритм.
Если $p = 2$ или $a = 0$, то $z = a$.
Иначе $p \ge 3$, $a \ge 1$.
Если $a^{(p-1)/2} \neq 1 \mod p$, то ответа не существует. Иначе же ответ есть.
Запускаем бесконечный цикл, пока не найдем ответ. В нем выбираем $i$ случайным образом из чисел $1, 2, \ldots, p - 1$. Считаем многочлен $T(x) := (x+i)^{(p-1)/2} - 1 \mod (x^2 - a) = bx + c$ бинарным возведением в степень и взятием по модулю каждый раз.
После чего если $b \neq 0$, то возьмем $z’ = c/b = c \cdot b^{p - 2} \mod p$ и проверим, подходит ли оно (возведем в квадрат). Если подходит, вернем, иначе продолжаем перебирать $i$.
Почему этот алгоритм работает? Крайние случаи очевидны. Иначе, числа, у которых есть корень, называются квадратичными вычетами. Давайте заметим, что у каждого числа ровно $2$ корня, либо ровно $0$ (если есть корень, то подходит и минус корень, а больше быть не может, потому что $z^2 = t^2 \mod p$ решается только так). То есть у нас есть ровно $(p-1)/2$ чисел, являющихся квадратичными вычетами и $(p-1)/2$, не являющихся. $x^{p-1} = 1 \mod p$, так что для проверки, является ли $a$ квадратом, можно возвести в $(p-1)/2$ степень. То есть проверка на несуществование ответа верная. Теперь осталось его найти!
Пусть $z^2 = a \mod p$. Корни многочлена $x^2 - a$ — это как раз $z$ и $-z$.
В каком случае они будут являться корнями многочлена $T$? Если $z + i$ и $-z + i$ являются квадратичными вычетами. Нас интересует ситуация, когда ровно один из них — квадратичный вычет. Пусть это $z + i$. Тогда $z$ — корень обоих многочленов, так что после взятия по модулю ответ делится на $x - z$. При этом ответ — многочлен не более, чем первой степени. Так что он имеет вид $dx - dz$. Тогда если мы поделим коэффициенты, то получим как раз $z$. Аналогично, если $-z$ — корень обоих многочленов. Обратите внимание, что $d \neq 0$, потому что в этом случае оба числа $\pm z$ являются корнями.
Осталось понять, почему этот алгоритм будет работать быстро. Итерация работает за $O(\log p)$. Докажем, что вероятность успеха — $1/2$, тогда матожидание количества шагов будет равно двум.
Если ровно одно из чисел $z + i$ и $-z + i$ является квадратичным вычетом, то $(z + i)^{(p-1)/2} \neq (-z + i)^{(p-1)/2} \mod p$. При этом одно из этих чисел — $1$, а другое — $-1$. Поделим одну часть на другую. $((z+i)/(-z+i))^{(p-1)/2} = -1 \mod p$. То есть это верно, если отношение — не вычет. Докажем, что все такие числа вида $(z+i)/(-z+i)$ различны для всех $i$, тогда это будет как раз биекция во все вычеты, так что вероятность равна $1/2$.
Пусть $(z +i)/(-z+i) = (z+j)/(-z+j) \mod p$. Домножим на знаменатели, сократим, получим, что $2iz = 2jz \mod p$, то есть $i = j \mod p$.
Вероятность не совсем $1/2$, потому что нужно, чтобы $z + i \neq 0$ и $-z + i \neq 0$. Так что нам не подходят еще два вычета. Алгоритм от этого не страдает, просто для малых $p$ вероятность становится чуть меньше. Можно доказать, что для $p = 3$ все равно все работает.