Discrete square root modulo

#	User	Rating
1	tourist	3757
2	jiangly	3647
3	Benq	3581
4	orzdevinwang	3570
5	Geothermal	3569
5	cnnfls_csy	3569
7	Radewoosh	3509
8	ecnerwala	3486
9	jqdai0815	3474
10	gyh20	3447

#	User	Contrib.
1	maomao90	171
2	awoo	165
3	adamant	163
4	TheScrasse	159
5	maroonrk	155
6	nor	154
7	-is-this-fft-	152
8	Petr	147
9	orz	146
10	pajenegod	145

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.

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