Efficient finding primitive roots modulo n using Python?

Question

I'm using the following code for finding primitive roots modulo n in Python:

Code:

def gcd(a,b):
    while b != 0:
        a, b = b, a % b
    return a

def primRoots(modulo):
    roots = []
    required_set = set(num for num in range (1, modulo) if gcd(num, modulo) == 1)

    for g in range(1, modulo):
        actual_set = set(pow(g, powers) % modulo for powers in range (1, modulo))
        if required_set == actual_set:
            roots.append(g)           
    return roots

if __name__ == "__main__":
    p = 17
    primitive_roots = primRoots(p)
    print(primitive_roots)

Output:

[3, 5, 6, 7, 10, 11, 12, 14]   

Code fragment extracted from: Diffie-Hellman (Github)

---

Can the primRoots method be simplified or optimized in terms of memory usage and performance/efficiency?

Answer 1

One quick change that you can make here (not efficiently optimum yet) is using list and set comprehensions:

def primRoots(modulo):
    coprime_set = {num for num in range(1, modulo) if gcd(num, modulo) == 1}
    return [g for g in range(1, modulo) if coprime_set == {pow(g, powers, modulo)
            for powers in range(1, modulo)}]

Now, one powerful and interesting algorithmic change that you can make here is to optimize your gcd function using memoization. Or even better you can simply use built-in gcd function form math module in Python-3.5+ or fractions module in former versions:

from functools import wraps
def cache_gcd(f):
    cache = {}

    @wraps(f)
    def wrapped(a, b):
        key = (a, b)
        try:
            result = cache[key]
        except KeyError:
            result = cache[key] = f(a, b)
        return result
    return wrapped

@cache_gcd
def gcd(a,b):
    while b != 0:
        a, b = b, a % b
    return a
# or just do the following (recommended)
# from math import gcd

Then:

def primRoots(modulo):
    coprime_set = {num for num in range(1, modulo) if gcd(num, modulo) == 1}
    return [g for g in range(1, modulo) if coprime_set == {pow(g, powers, modulo)
            for powers in range(1, modulo)}]

As mentioned in comments, as a more pythoinc optimizer way you can use fractions.gcd (or for Python-3.5+ math.gcd).

Answer 2

Based on the comment of Pete and answer of Kasramvd, I can suggest this:

from math import gcd as bltin_gcd

def primRoots(modulo):
    required_set = {num for num in range(1, modulo) if bltin_gcd(num, modulo) }
    return [g for g in range(1, modulo) if required_set == {pow(g, powers, modulo)
            for powers in range(1, modulo)}]

print(primRoots(17))

Output:

[3, 5, 6, 7, 10, 11, 12, 14]

---

Changes:

• It now uses pow method's 3-rd argument for the modulo.
• Switched to gcd built-in function that's defined in math (for Python 3.5) for a speed boost.

---

Additional info about built-in gcd is here: Co-primes checking

Answer 3

In the special case that p is prime, the following is a good bit faster:

import sys

# translated to Python from http://www.bluetulip.org/2014/programs/primitive.js
# (some rights may remain with the author of the above javascript code)

def isNotPrime(possible):
    # We only test this here to protect people who copy and paste
    # the code without reading the first sentence of the answer.
    # In an application where you know the numbers are prime you
    # will remove this function (and the call). If you need to
    # test for primality, look for a more efficient algorithm, see
    # for example Joseph F's answer on this page.
    i = 2
    while i*i <= possible:
        if (possible % i) == 0:
            return True
        i = i + 1
    return False

def primRoots(theNum):
    if isNotPrime(theNum):
        raise ValueError("Sorry, the number must be prime.")
    o = 1
    roots = []
    r = 2
    while r < theNum:
        k = pow(r, o, theNum)
        while (k > 1):
            o = o + 1
            k = (k * r) % theNum
        if o == (theNum - 1):
            roots.append(r)
        o = 1
        r = r + 1
    return roots

print(primRoots(int(sys.argv[1])))