Miller-Rabin Primality Test

miller-rabin.py

from random import randrange


def miller_rabin(n, k):
    if n == 2 or n == 3:
        return True
    if n < 2 or ~n & 1:
        return False
    s, d = 0, n - 1
    while ~d & 1:
        d, s = d // 2, s + 1
    for i in range(k):
        c = False
        r = randrange(2, n - 1)
        x = pow(r, d, n)
        if x == 1 or x == n - 1:
            continue
        for j in range(s - 1):
            x = (x * x) % n
            if x == 1:
                return False
            if x == n - 1:
                c = True
                break
        if not c:
            return False
    return True


def main():
    N, K = int(argv[1]), 8
    if len(argv) is 3:
        K = int(argv[2])
    if miller_rabin(N, K):
        print(N, 'may be prime')
    else:
        print(N, 'is not prime')

if __name__ == '__main__':
    from sys import argv
    main()