Number Theory & Algebra contents

Discrete Root

Solve x^k ≡ a (mod n) by moving to exponents with a primitive root and a linear congruence.

Advanced4 min readdiscrete rootprimitive rootmodular arithmeticnth roots

Read first: Discrete Logarithm, Primitive Roots, Linear Congruence Equation

The discrete root problem: given $n$ , $k$ and $a$ , find all $x$ with

x^k \equiv a \pmod n

For $k = 2$ these are modular square roots. We treat the case where $n$ has a primitive root $g$ (for example $n$ prime) and $a \not\equiv 0$ .

Algorithm

Every non-zero residue is a power of $g$ . Write $x = g^y$ and $a = g^{\text{ind}(a)}$ , where $\text{ind}(a)$ is the discrete logarithm of $a$ to the base $g$ . The equation becomes

g^{yk} \equiv g^{\text{ind}(a)} \pmod n \iff k\,y \equiv \text{ind}(a) \pmod{\varphi(n)}

because $g$ has order $\varphi(n)$ . That's a linear congruence in $y$ . Let $d = \gcd(k, \varphi(n))$ :

if $d \nmid \text{ind}(a)$ there is no root;
otherwise there are $d$ solutions $y$ , which yield $d$ different values of $x = g^y$ .

python

from math import gcd, isqrt

def prime_factors(n):
    ps, d = [], 2
    while d * d <= n:
        if n % d == 0:
            ps.append(d)
            while n % d == 0:
                n //= d
        d += 1
    if n > 1:
        ps.append(n)
    return ps

def phi(n):
    r = n
    for p in prime_factors(n):
        r -= r // p
    return r

def primitive_root_of_prime(p):
    if p == 2:
        return 1
    order = p - 1
    qs = prime_factors(order)
    return next(g for g in range(2, p) if all(pow(g, order // q, p) != 1 for q in qs))

def dlog(g, a, p):
    """Baby-step giant-step for prime p: x with g^x = a (mod p)."""
    n = isqrt(p) + 1
    baby = {}
    cur = a % p
    for q in range(n + 1):
        baby[cur] = q
        cur = cur * g % p
    step = pow(g, n, p)
    cur = 1
    for i in range(1, n + 1):
        cur = cur * step % p
        if cur in baby:
            return n * i - baby[cur]
    return None

def discrete_roots(k, a, p):
    """All x in [0, p) with x^k = a (mod p), p prime."""
    a %= p
    if a == 0:
        return [0]
    g = primitive_root_of_prime(p)
    t = dlog(g, a, p)                       # index of a
    order = p - 1
    d = gcd(k, order)
    if t % d:
        return []
    k1, t1, o1 = k // d, t // d, order // d
    y0 = t1 * pow(k1, -1, o1) % o1 if o1 > 1 else 0
    return sorted(pow(g, y0 + j * o1, p) for j in range(d))

assert discrete_roots(2, 4, 7) == [2, 5]                 # 2^2 = 4, 5^2 = 25 = 4 (mod 7)
assert discrete_roots(2, 3, 7) == []                     # 3 is not a square mod 7
assert discrete_roots(3, 1, 7) == [1, 2, 4]              # cube roots of unity
assert discrete_roots(3, 6, 7) == [3, 5, 6]              # three cube roots of 6
assert discrete_roots(5, 6, 11) == []                    # fifth powers mod 11 are only 1 and 10

for p in (2, 3, 5, 7, 11, 13, 17, 19, 23, 101):
    for k in range(1, 15):
        for a in range(p):
            expected = [x for x in range(p) if pow(x, k, p) == a]
            assert discrete_roots(k, a, p) == expected, (k, a, p)

The last block verifies every $(k, a)$ for several primes against brute force.

Finding all solutions from one

If $x_0$ is one root, all roots are $x_0$ times the $k$ -th roots of unity (solutions of $z^k \equiv 1$ ), which form a cyclic group of size $\gcd(k, \varphi(n))$ generated by $g^{\varphi(n)/d}$ . That's what the code exploits by stepping $y$ by $\varphi(n)/d$ .

Application: modular square roots

For $k = 2$ , a non-zero $a$ has either 0 or 2 roots modulo an odd prime, and $a$ is a quadratic residue exactly when $a^{(p-1)/2} \equiv 1$ (Euler's criterion):

python

def is_quadratic_residue(a, p):
    return pow(a, (p - 1) // 2, p) == 1

for p in (7, 11, 13, 101):
    for a in range(1, p):
        assert is_quadratic_residue(a, p) == bool(discrete_roots(2, a, p))

For very large primes there are faster special algorithms (Tonelli-Shanks, Cipolla) that avoid the discrete logarithm entirely.

Practice problems

Codeforces - Lunar New Year and a Recursive Sequence