Discrete square root modulo

#	User	Rating
1	tourist	4009
2	jiangly	3823
3	Benq	3738
4	Radewoosh	3633
5	jqdai0815	3620
6	orzdevinwang	3529
7	ecnerwala	3446
8	Um_nik	3396
9	ksun48	3390
10	gamegame	3386

#	User	Contrib.
1	cry	167
2	Um_nik	163
3	maomao90	162
3	atcoder_official	162
5	adamant	159
6	-is-this-fft-	158
7	awoo	157
8	TheScrasse	154
9	Dominater069	153
9	nor	153

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