Discrete square root modulo

#	User	Rating
1	tourist	3985
2	jiangly	3814
3	jqdai0815	3682
4	Benq	3529
5	orzdevinwang	3526
6	ksun48	3517
7	Radewoosh	3410
8	hos.lyric	3399
9	ecnerwala	3392
9	Um_nik	3392

#	User	Contrib.
1	cry	169
2	maomao90	162
2	Um_nik	162
4	atcoder_official	161
5	djm03178	158
6	-is-this-fft-	157
7	adamant	155
8	awoo	154
8	Dominater069	154
10	luogu_official	150

Just recently I had a need to solve the following problem:

Given $$$a$$$ and $$$m$$$, find all $$$x$$$ such that: $$$x^2 \equiv a \pmod{m}$$$.

After some time searching the Internet and carefully analyzing each case, I got a solution that works for any $$$a$$$ and $$$m > 0$$$ in $$$O(\sqrt{m})$$$ time, or in $$$O(\log^2(m))$$$ time if $$$a$$$ and $$$m$$$ are coprime (if there is a factorization of $$$m$$$).

Below are the cases I considered and links to articles (or discussions) for algorithms:

$$$x^2 \equiv a \pmod{p}$$$, $$$a \ne 0$$$, $$$p$$$ — odd prime (link).
$$$x^2 \equiv a \pmod{p^k}$$$, $$$a \ne 0$$$, $$$k > 0$$$, $$$\gcd(a, p) = 1$$$, $$$p$$$ — odd prime (link).
$$$x^2 \equiv a \pmod{2^k}$$$, $$$a \ne 0$$$, $$$k > 0$$$, $$$\gcd(a, 2) = 1$$$ (link).
$$$x^2 \equiv 0 \pmod{p^k}$$$, $$$k > 0$$$, $$$p$$$ — prime. Note that in this case $$$x^2$$$ must be divisible by $$$p^k$$$. Then $$$x$$$ must be divisible by $$$p^{\lceil\frac{k}{2}\rceil}$$$. From this we obtain that $$$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$$$ — prime. Solutions exist only for $$$s = 2 \cdot t$$$. Proved by contradiction. Then it is obvious that $$$p^t$$$ is one of the solutions ($$$(p^t)^2 = p^s$$$). Let's find the remaining solutions. Let $$$x = p^t + y$$$, then $$$x^2 = p^s + 2 \cdot p^t \cdot y + y^2$$$. Note that $$$2 \cdot p^t \cdot y + y^2$$$ must be divisible by $$$p^k$$$. In this case, $$$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$$$ — prime. Let $$$x_1$$$ be the solution for $$$x^2 \equiv q \pmod{p^k}$$$, and let $$$x_2$$$ be the solution for $$$x^2 \equiv p^s \pmod{p^k}$$$. Then $$$x = x_1 \cdot x_2$$$ will be a solution to the original problem.

The last three cases had to be deduced independently, so there are no links. I was too lazy to document the proofs in this post, although in general they are not very complicated.

The general solution is found by using the Chinese Remainder Theorem and reducing to one of the above cases.

I will also give a complete implementation in Python (a function that solves the general case — discrete_sqrt):

Implementation

def _bin_pow(a, b, m):
    c = 1
    while b > 0:
        if b % 2 == 1:
            c = c * a % m
        a = a * a % m
        b //= 2
    return c


# x^2 = a (mod p)
def _tonelli_shanks(a, p):
    assert a % p != 0
    assert p > 2
    q = p - 1
    s = 0
    while q % 2 == 0:
        q //= 2
        s += 1
    assert s > 0
    if _bin_pow(a, (p - 1) // 2, p) != 1:
        return []
    z = 2
    while _bin_pow(z, (p - 1) // 2, p) != p - 1:
        z += 1
        assert z < p
    c = _bin_pow(z, q, p)
    r = _bin_pow(a, (q + 1) // 2, p)
    t = _bin_pow(a, q, p)
    while t != 1:
        i = 1
        while i < s and _bin_pow(t, 2 ** i, p) != 1:
            i += 1
        assert i < s
        b = _bin_pow(c, 2 ** (s - i - 1), p)
        r = r * b % p
        t = t * b * b % p
        c = b * b % p
        s = i
    return sorted([r, p - r])


# x^2 = a (mod p^k)
def _generalized_tonelli_shanks(a, p, k):
    assert a % p != 0
    assert p > 2
    assert k > 0
    ans = []
    m = p ** k
    d = _bin_pow(a, (m - 2 * m // p + 1) // 2, m)
    for x in _tonelli_shanks(a, p):
        ans.append(_bin_pow(x, m // p, m) * d % m)
    return sorted(ans)


# x^2 = a (mod 2^k)
def _discrete_sqrt_2(a, k):
    assert k > 0
    assert a > 0
    if k == 1:
        return [1]
    if a % 8 != 1:
        return []
    ans = set()
    for x in [1, 3]:
        for i in range(3, k):
            j = ((x * x - a) // (2 ** i) % 2 + 2) % 2
            x = x + j * 2 ** (i - 1)
        ans.add(x)
        ans.add(2 ** k - x)
    return sorted(list(ans))


# x^2 = 0 (mod p^k)
def _discrete_sqrt_zero(p, k):
    assert p > 1
    assert k > 0
    ans = [0]
    x = p ** ((k + 1) // 2)
    for i in range(1, (p ** k) // x):
        ans.append(x * i)
    return ans


# x^2 = a (mod p^k)
def _discrete_sqrt_prime(a, p, k):
    assert p > 1
    assert k > 0
    if a == 0:
        return _discrete_sqrt_zero(p, k)
    q, s = a, 0
    while q % p == 0:
        q //= p
        s += 1
    if s == 0:
        if p == 2:
            return _discrete_sqrt_2(a, k)
        return _generalized_tonelli_shanks(a, p, k)
    if s % 2 != 0:
        return []
    xq = _discrete_sqrt_prime(q, p, k)
    x0 = p ** (s // 2)
    d = p ** max((k + 1) // 2, k - s // 2 - (1 if p == 2 else 0))
    ans = set()
    for i in range((p ** k) // d):
        for x in xq:
            ans.add((x0 + i * d) * x % (p ** k))
    return sorted(list(ans))


def _ex_gcd(a, b):
    x, xp = 1, 0
    y, yp = 0, 1
    while b != 0:
        q = a // b
        a, b = b, a - q * b
        x, xp = xp, x - q * xp
        y, yp = yp, y - q * yp
    return (a, x, y)


def _factorize(n):
    assert n > 0
    ans = []
    i = 2
    while i * i <= n:
        if n % i != 0:
            i += 1
            continue
        ans.append(i)
        n //= i
    if n > 1:
        ans.append(n)
    return ans


# x^2 = a (mod m)
def discrete_sqrt(a, m):
    assert m > 0
    ans = [0]
    mans = 1
    fm = _factorize(m)
    it = 0
    while it < len(fm):
        jt = it
        while jt < len(fm) and fm[jt] == fm[it]:
            jt += 1
        p = fm[it]
        k = jt - it
        pk = p ** k
        x = _discrete_sqrt_prime(a % pk, p, k)
        g, x1, _ = _ex_gcd(mans, pk)
        assert g == 1
        nans = []
        for r0 in ans:
            for r1 in x:
                nans.append(((r0 + x1 * (r1 - r0) % pk * mans) % (mans * pk) + (mans * pk)) % (mans * pk))
        ans = nans
        mans *= pk
        it = jt
    return sorted(ans)

wilcot's blog