number_theory_functions.py

from random import randrange
from math import floor
from numpy import binary_repr


def extended_gcd(a, b):
    """
    Returns the extended gcd of a and b

    Parameters
    ----------
    a : Input data.
    b : Input data.
    Returns
    -------
    (d, x, y): d = gcd(a,b) = a*x + b*y
    """

    # if initially b>a, after one iteration they will switch places
    # and the calculation can go on normally.

    if a == 0:
        return b, 0, 1

    d, temp_a, temp_b = extended_gcd(b % a, a)

    a = temp_b - floor(b / a) * temp_a
    b = temp_a

    return d, a, b


def modular_inverse(a, n):
    """
    Returns the inverse of a modulo n if one exists

    Parameters
    ----------
    a : Input data.
    n : Input data.

    Returns
    -------
    x: such that (a*x % n) == 1 and 0 <= x < n if one exists, else None
    """
    gcd, x, y = extended_gcd(a, n)
    if gcd != 1:
        return None
    return x % n


def modular_exponent(a, d, n):
    """
    Returns a to the power of d modulo n

    Parameters
    ----------
    a : The exponential's base.
    d : The exponential's exponent.
    n : The exponential's modulus.

    Returns
    -------
    b: such that b == (a**d) % n
    """
    binary = binary_repr(d)
    c = 1
    prev = (1, a)  # a ** 1 % n = a
    for index, i in enumerate(reversed(list(binary))):
        if int(i) != 0:
            power = 2 ** index
            p = int(power / prev[0])
            num = prev[1] ** p
            num = num % n
            prev = (power, num)  # a ** power % n = num
            c = (c * num) % n
    return c


def miller_rabin(n):
    """
    Checks the primality of n using the Miller-Rabin test

    Parameters
    ----------
    n : The number to check

    Returns
    -------
    b: If n is prime, b is guaranteed to be True.
    If n is not a prime, b has a 3/4 chance at least to be False
    """
    a = randrange(1, n)
    k = 0
    d = n - 1
    while d % 2 == 0:
        k = k + 1
        d = d // 2
    x = modular_exponent(a, d, n)
    if x == 1 or x == n - 1:
        return True
    for _ in range(k):
        x = (x * x) % n
        if x == 1:
            return False
        if x == n - 1:
            return True
    return False


def is_prime(n):
    """
    Checks the primality of n

    Parameters
    ----------
    n : The number to check

    Returns
    -------
    b: If n is prime, b is guaranteed to be True.
    If n is not a prime, b has a chance of less than 1e-10 to be True
    """
    for _ in range(10):
        if not miller_rabin(n):
            return False
    return True


def generate_prime(digits):
    for i in range(digits * 20):
        n = randrange(10 ** (digits - 1), 10 ** digits)
        if is_prime(n):
            return n
    return None