pohlig-hellman.py

import math
import collections

def is_prime(n):
    if n < 2:
        return False
    if n == 2:
        return True
    if n % 2 == 0:
        return False
    for i in range(3, int(math.sqrt(n)) + 1, 2):
        if n % i == 0:
            return False
    return True

def genprimes():
    i = 2
    while True:
        if is_prime(i):
            yield i
        i += 1

def extended_gcd(a, b):
    if b == 0:
        return (a, 1, 0)
    else:
        d, x, y = extended_gcd(b, a % b)
        return (d, y, x - (a // b) * y)   

def chinese_remainder(pairs):
    '''Return the solution to the Chinese Remainder Theorem, (x, M)
    Pairs contains tuples (a, m) with all m's positive and coprime.
    Return the smallest nonnegative integer x so 
    x mod m  = a mod m for each (a, m) in pairs.
    M is the product ofthe m's.
    '''
    (a, m)=pairs[0]
    for (b,p) in pairs[1:]:
        k=((b-a)*extended_gcd(m,p)[1]) % p #moduli coprime so inverse exists for m mod p
        a=(a+m*k) % (m*p)# joining a mod m and b mod p gives a mod(mp)
        m *= p # mod mp
    return (a,m)
 

def factorize(n):
    factors = collections.defaultdict(int)
    primes = genprimes()

    while n > 1:
        p = next(primes)
        while n % p == 0:
            n = n // p
            factors[p] += 1
    return factors


def pohlig_hellman(a, b, p):
    '''
    solve a^x = b (mod p)
    follows https://en.wikipedia.org/wiki/Pohlig%E2%80%93Hellman_algorithm#The_general_algorithm
    seems to work for p = prime powers but not other all other cases
    '''
    factors = factorize(p)
    #print(factors)

    result = 0
    congruences = []
    for q, e in factors.items():

        g = pow(a, (p) // pow(q, e), p)
        h = pow(b, (p) // pow(q, e), p)

        z = None
        for xi in range(0, pow(q, e)):
            if pow(g, xi, p) == h:
                z = xi
                break

        # if z is None:
        #     print("NOT FOUND")

        congruences.append((z, pow(q, e)))
 
        #: ', q, e, g, h, z)
        
    sol = chinese_remainder(congruences)
    return sol[0]



if __name__ == '__main__':


    # Works for any p which is a prime power!
    a = 40
    b = 45
    p = 47
    x = pohlig_hellman(a, b, p) # x = 25
    print("SOLUTION:", x)

    # Doesn't work
    # a = 39 
    # b = 49 
    # p = 74 
    # x = pohlig_hellman(a, b, p) # x = 28
    # print("SOLUTION:", x)

    # Doesn't work
    # a = 19 
    # b = 423 
    # p = 78 
    # x = pohlig_hellman(a, b, p) # x = 275
    # print("SOLUTION:", x)

    # Doesn't work
    # a = 71 
    # b = 65 
    # p = 86 
    # x = pohlig_hellman(a, b, p) # x = 45
    # print("SOLUTION:", x)