Pollard's Rho algorithm

pollards_rho.py

from argparse import ArgumentParser
import numpy as np
gcd_table = {}

def gcd(a, b):

    while a % b != 0:
        a, b = b, a % b
    return b
bound = 10000

def b_gcd(u, v):
  """
  binary gcd (not mine)
  """
    if u in gcd_table:
        table = gcd_table.get(u)
        if v in table:
            return table[v]
    elif u < bound:
        gcd_table[u] = {}

    if u == v:
        return u
    elif u == 0:
        return v
    elif v == 0:
        return u
    # u is even
    elif u & 1 == 0:
        # v is even
        if v & 1 == 0:
            b = 2*gcd(u >> 1, v >> 1)
            if u < bound: gcd_table[u][v] = b
            return b
        # v is odd
        else:
            b = gcd(u >> 1, v)
            if u < bound: gcd_table[u][v] = b
            return b
    # u is odd
    elif u & 1 != 0:
        # v is even
        if v & 1 == 0:
            b = gcd(u, v >> 1)
            if u < bound: gcd_table[u][v] = b
            return b
        # v is odd and u is greater than v
        elif u > v and v & 1 != 0:
            b = gcd((u-v) >> 1, v)
            if u < bound: gcd_table[u][v] = b
            return b
        # v is odd and u is smaller than v
        else:
            b = gcd((v-u) >> 1, u)
            if u < bound: gcd_table[u][v] = b
            return b


def pollards_rho(n, g, x=2, brent=False, output=True, rand=False):
    if rand:
        n_64 = min(n, 2**64-1)
        a = np.random.randint(1, n_64-3)
        b = np.random.randint(1, n_64-1)
        print(a, b)
        x = a
        y = b
    else:
        y = x
    cycle_size = 2
    d = 1
    latex = output
    iteration = 0
    if latex:
        print("\\hline \n %d & %d & %d & %d \\\\" % (iteration, x, y, d))

    while d == 1:

        if brent:
            count = 1
            while count <= cycle_size and d <= 1:
                x = g(x) % n
                d = b_gcd(abs(x - y), n)
                count += 1
            cycle_size *= 2
            y = x
        else:
            x = g(x) % n
            y = g(g(y)) % n
            d = b_gcd(abs(x - y), n)
            # count += 1
            # cycle_size *= 2

        iteration += 1
        if latex:
            print("\\hline \n %d & %d & %d & %d \\\\" % (iteration, x, y, d))
        elif output:
            print("At iteration %d" % iteration)
            print("x = %d" % x)
            print("y = %d" % y)
            print("d = %d" % d)
            print()

        if iteration % 100 == 0:
            print(iteration)


    if d == n:
        print("Unsuccessful")
        return False

    return d, iteration


if __name__ == "__main__":
    ap = ArgumentParser()
    ap.add_argument("-n", dest="n", type=int, default=10403)
    ap.add_argument("-x", dest="x", type=int, default=2)
    ap.add_argument("-o", dest="output", type=bool, default=False)
    ap.add_argument("-r", "--rand", dest="rand", type=bool, default=False)
    ap.add_argument("-c", "--cycle_alg", type=str,
                    choices=["brent", "floyd"], default="floyd")
    args = ap.parse_args()

    result = pollards_rho(args.n, lambda x: (x*x + 1), x=args.x,
          brent=(args.cycle_alg == "brent"), output=args.output, rand=args.rand)
    print(result)