From 2a8f0f0747650050a5e8ab364e2acb005b085cea Mon Sep 17 00:00:00 2001 From: Clear20-22 Date: Tue, 1 Sep 2026 13:16:55 +0600 Subject: [PATCH 1/2] Add Tonelli-Shanks algorithm for modular square roots --- maths/tonelli_shanks.py | 142 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 142 insertions(+) create mode 100644 maths/tonelli_shanks.py diff --git a/maths/tonelli_shanks.py b/maths/tonelli_shanks.py new file mode 100644 index 000000000000..d809e882e713 --- /dev/null +++ b/maths/tonelli_shanks.py @@ -0,0 +1,142 @@ +""" +Tonelli-Shanks Algorithm for Modular Square Roots. + +Reference: https://en.wikipedia.org/wiki/Tonelli%E2%80%93Shanks_algorithm +Reference: https://cp-algorithms.com/algebra/tonelli-shanks.html + +Given an integer n and an odd prime p, the Tonelli-Shanks algorithm computes +an integer x such that: + x^2 = n (mod p) +If n is a quadratic residue modulo p, it returns two solutions (r1, r2) +where r1 <= r2 and r1 + r2 = p (or (0, 0) if n = 0 mod p). +""" + +from __future__ import annotations + + +def legendre_symbol(n: int, p: int) -> int: + """ + Compute the Legendre Symbol (n / p) modulo an odd prime p. + + Euler's Criterion states that: + (n / p) = n^((p - 1) / 2) (mod p) + + Returns: + 1 if n is a quadratic residue modulo p and n != 0 (mod p) + -1 if n is a quadratic non-residue modulo p + 0 if n = 0 (mod p) + + >>> legendre_symbol(5, 11) + 1 + >>> legendre_symbol(2, 7) + 1 + >>> legendre_symbol(3, 7) + -1 + >>> legendre_symbol(0, 7) + 0 + >>> legendre_symbol(14, 7) + 0 + """ + ls = pow(n % p, (p - 1) // 2, p) + return ls if ls <= 1 else -1 + + +def tonelli_shanks(n: int, p: int) -> tuple[int, int]: + """ + Find solutions to x^2 = n (mod p) for an odd prime p using Tonelli-Shanks algorithm. + + Time Complexity: O(log^2 p) on average. + + Parameters: + n: The integer whose square root modulo p is to be found. + p: An odd prime modulus. + + Returns: + A tuple (r1, r2) containing the two square roots modulo p such that r1 <= r2. + + Raises: + ValueError: If p is not an odd prime >= 3. + ValueError: If n is not a quadratic residue modulo p. + + >>> tonelli_shanks(5, 11) + (4, 7) + >>> tonelli_shanks(10, 13) + (6, 7) + >>> tonelli_shanks(0, 7) + (0, 0) + >>> tonelli_shanks(2, 7) + (3, 4) + >>> tonelli_shanks(28, 7) + (0, 0) + >>> tonelli_shanks(3, 7) + Traceback (most recent call last): + ... + ValueError: 3 is not a quadratic residue modulo 7. + >>> tonelli_shanks(5, 4) + Traceback (most recent call last): + ... + ValueError: Modulus p must be an odd prime (got 4). + >>> tonelli_shanks(5, 1) + Traceback (most recent call last): + ... + ValueError: Modulus p must be an odd prime (got 1). + """ + if p <= 2 or p % 2 == 0: + msg = f"Modulus p must be an odd prime (got {p})." + raise ValueError(msg) + + n = n % p + if n == 0: + return 0, 0 + + if legendre_symbol(n, p) != 1: + msg = f"{n} is not a quadratic residue modulo {p}." + raise ValueError(msg) + + # Factor out powers of 2 from p - 1: p - 1 = q * 2^s, with q odd + q = p - 1 + s = 0 + while q % 2 == 0: + q //= 2 + s += 1 + + # Case 1: p = 3 (mod 4), s = 1 + if s == 1: + root = pow(n, (p + 1) // 4, p) + return min(root, p - root), max(root, p - root) + + # Case 2: Search for a quadratic non-residue z modulo p + z = 2 + while legendre_symbol(z, p) != -1: + z += 1 + + c = pow(z, q, p) + x = pow(n, (q + 1) // 2, p) + t = pow(n, q, p) + m = s + + while t != 1: + # Find the smallest i (0 < i < m) such that t^(2^i) = 1 (mod p) + i = 0 + t2i = t + while t2i != 1 and i < m: + t2i = pow(t2i, 2, p) + i += 1 + + if i == m: + msg = f"Failed to find square root for {n} mod {p}." + raise ValueError(msg) + + b = pow(c, 1 << (m - i - 1), p) + x = (x * b) % p + t = (t * b * b) % p + c = (b * b) % p + m = i + + return min(x, p - x), max(x, p - x) + + +if __name__ == "__main__": + import doctest + + doctest.testmod() From 8a8bbdd371c6c66ddf2156feb14f46ca7b67de2e Mon Sep 17 00:00:00 2001 From: Clear20-22 Date: Tue, 1 Sep 2026 13:22:10 +0600 Subject: [PATCH 2/2] Refactor: use descriptive variable and parameter names in tonelli_shanks --- maths/tonelli_shanks.py | 133 ++++++++++++++++++++-------------------- 1 file changed, 68 insertions(+), 65 deletions(-) diff --git a/maths/tonelli_shanks.py b/maths/tonelli_shanks.py index d809e882e713..a9ec83d9db4a 100644 --- a/maths/tonelli_shanks.py +++ b/maths/tonelli_shanks.py @@ -4,27 +4,27 @@ Reference: https://en.wikipedia.org/wiki/Tonelli%E2%80%93Shanks_algorithm Reference: https://cp-algorithms.com/algebra/tonelli-shanks.html -Given an integer n and an odd prime p, the Tonelli-Shanks algorithm computes -an integer x such that: - x^2 = n (mod p) -If n is a quadratic residue modulo p, it returns two solutions (r1, r2) -where r1 <= r2 and r1 + r2 = p (or (0, 0) if n = 0 mod p). +Given an integer residue and an odd prime modulus, the Tonelli-Shanks algorithm +computes an integer root such that: + root^2 = residue (mod modulus) +If residue is a quadratic residue modulo modulus, it returns two solutions (r1, r2) +where r1 <= r2 and r1 + r2 = modulus (or (0, 0) if residue = 0 mod modulus). """ from __future__ import annotations -def legendre_symbol(n: int, p: int) -> int: +def legendre_symbol(number: int, modulus: int) -> int: """ - Compute the Legendre Symbol (n / p) modulo an odd prime p. + Compute the Legendre Symbol (number / modulus) modulo an odd prime modulus. Euler's Criterion states that: - (n / p) = n^((p - 1) / 2) (mod p) + (number / modulus) = number^((modulus - 1) / 2) (mod modulus) Returns: - 1 if n is a quadratic residue modulo p and n != 0 (mod p) - -1 if n is a quadratic non-residue modulo p - 0 if n = 0 (mod p) + 1 if number is a quadratic residue modulo modulus and number != 0 (mod modulus) + -1 if number is a quadratic non-residue modulo modulus + 0 if number = 0 (mod modulus) >>> legendre_symbol(5, 11) 1 @@ -37,26 +37,27 @@ def legendre_symbol(n: int, p: int) -> int: >>> legendre_symbol(14, 7) 0 """ - ls = pow(n % p, (p - 1) // 2, p) - return ls if ls <= 1 else -1 + symbol = pow(number % modulus, (modulus - 1) // 2, modulus) + return symbol if symbol <= 1 else -1 -def tonelli_shanks(n: int, p: int) -> tuple[int, int]: +def tonelli_shanks(residue: int, modulus: int) -> tuple[int, int]: """ - Find solutions to x^2 = n (mod p) for an odd prime p using Tonelli-Shanks algorithm. + Find solutions to root^2 = residue (mod modulus) for an odd prime modulus + using the Tonelli-Shanks algorithm. - Time Complexity: O(log^2 p) on average. + Time Complexity: O(log^2 modulus) on average. Parameters: - n: The integer whose square root modulo p is to be found. - p: An odd prime modulus. + residue: The integer whose square root modulo modulus is to be found. + modulus: An odd prime modulus. Returns: - A tuple (r1, r2) containing the two square roots modulo p such that r1 <= r2. + A tuple (r1, r2) containing the two square roots modulo modulus (r1 <= r2). Raises: - ValueError: If p is not an odd prime >= 3. - ValueError: If n is not a quadratic residue modulo p. + ValueError: If modulus is not an odd prime >= 3. + ValueError: If residue is not a quadratic residue modulo modulus. >>> tonelli_shanks(5, 11) (4, 7) @@ -81,59 +82,61 @@ def tonelli_shanks(n: int, p: int) -> tuple[int, int]: ... ValueError: Modulus p must be an odd prime (got 1). """ - if p <= 2 or p % 2 == 0: - msg = f"Modulus p must be an odd prime (got {p})." + if modulus <= 2 or modulus % 2 == 0: + msg = f"Modulus p must be an odd prime (got {modulus})." raise ValueError(msg) - n = n % p - if n == 0: + residue = residue % modulus + if residue == 0: return 0, 0 - if legendre_symbol(n, p) != 1: - msg = f"{n} is not a quadratic residue modulo {p}." + if legendre_symbol(residue, modulus) != 1: + msg = f"{residue} is not a quadratic residue modulo {modulus}." raise ValueError(msg) - # Factor out powers of 2 from p - 1: p - 1 = q * 2^s, with q odd - q = p - 1 - s = 0 - while q % 2 == 0: - q //= 2 - s += 1 - - # Case 1: p = 3 (mod 4), s = 1 - if s == 1: - root = pow(n, (p + 1) // 4, p) - return min(root, p - root), max(root, p - root) - - # Case 2: Search for a quadratic non-residue z modulo p - z = 2 - while legendre_symbol(z, p) != -1: - z += 1 - - c = pow(z, q, p) - x = pow(n, (q + 1) // 2, p) - t = pow(n, q, p) - m = s - - while t != 1: - # Find the smallest i (0 < i < m) such that t^(2^i) = 1 (mod p) - i = 0 - t2i = t - while t2i != 1 and i < m: - t2i = pow(t2i, 2, p) - i += 1 - - if i == m: - msg = f"Failed to find square root for {n} mod {p}." + # Factor out powers of 2 from modulus - 1: modulus - 1 = odd_factor * 2^two_power + odd_factor = modulus - 1 + two_power = 0 + while odd_factor % 2 == 0: + odd_factor //= 2 + two_power += 1 + + # Case 1: modulus = 3 (mod 4), two_power = 1 + if two_power == 1: + root = pow(residue, (modulus + 1) // 4, modulus) + return min(root, modulus - root), max(root, modulus - root) + + # Case 2: Search for a quadratic non-residue non_residue modulo modulus + non_residue = 2 + while legendre_symbol(non_residue, modulus) != -1: + non_residue += 1 + + multiplier = pow(non_residue, odd_factor, modulus) + root = pow(residue, (odd_factor + 1) // 2, modulus) + reduced_residue = pow(residue, odd_factor, modulus) + exponent = two_power + + while reduced_residue != 1: + # Find the smallest power_step (0 < power_step < exponent) + # such that reduced_residue^(2^power_step) = 1 (mod modulus) + power_step = 0 + temp_residue = reduced_residue + while temp_residue != 1 and power_step < exponent: + temp_residue = pow(temp_residue, 2, modulus) + power_step += 1 + + if power_step == exponent: + msg = f"Failed to find square root for {residue} mod {modulus}." raise ValueError(msg) - b = pow(c, 1 << (m - i - 1), p) - x = (x * b) % p - t = (t * b * b) % p - c = (b * b) % p - m = i + step_multiplier = pow(multiplier, 1 << (exponent - power_step - 1), modulus) + step_multiplier_squared = (step_multiplier * step_multiplier) % modulus + root = (root * step_multiplier) % modulus + reduced_residue = (reduced_residue * step_multiplier_squared) % modulus + multiplier = step_multiplier_squared + exponent = power_step - return min(x, p - x), max(x, p - x) + return min(root, modulus - root), max(root, modulus - root) if __name__ == "__main__":