diff --git a/maths/tonelli_shanks.py b/maths/tonelli_shanks.py new file mode 100644 index 000000000000..a9ec83d9db4a --- /dev/null +++ b/maths/tonelli_shanks.py @@ -0,0 +1,145 @@ +""" +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 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(number: int, modulus: int) -> int: + """ + Compute the Legendre Symbol (number / modulus) modulo an odd prime modulus. + + Euler's Criterion states that: + (number / modulus) = number^((modulus - 1) / 2) (mod modulus) + + Returns: + 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 + >>> legendre_symbol(2, 7) + 1 + >>> legendre_symbol(3, 7) + -1 + >>> legendre_symbol(0, 7) + 0 + >>> legendre_symbol(14, 7) + 0 + """ + symbol = pow(number % modulus, (modulus - 1) // 2, modulus) + return symbol if symbol <= 1 else -1 + + +def tonelli_shanks(residue: int, modulus: int) -> tuple[int, int]: + """ + Find solutions to root^2 = residue (mod modulus) for an odd prime modulus + using the Tonelli-Shanks algorithm. + + Time Complexity: O(log^2 modulus) on average. + + Parameters: + 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 modulus (r1 <= r2). + + Raises: + 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) + >>> 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 modulus <= 2 or modulus % 2 == 0: + msg = f"Modulus p must be an odd prime (got {modulus})." + raise ValueError(msg) + + residue = residue % modulus + if residue == 0: + return 0, 0 + + 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 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) + + 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(root, modulus - root), max(root, modulus - root) + + +if __name__ == "__main__": + import doctest + + doctest.testmod()