《编程思维与实践》大作业之质数序列迭代器和生成器

本次大作业要求写一个迭代器和一个生成器来生成质数序列。
下面使用轮子优化的埃氏筛实现迭代器和Miller-Rabin素性测试算法实现生成器。前者可以生成2开始的特定范围质数序列，后者可以生成特定数值开始的质数序列。

import math
from itertools import islice

class PrimeSeq:
    """
    An iterable class to generate a prime sequence from 2 to max_num
    """
    def _sieve(self, base, max_num, is_prime):
        for i in range(base * base, self.max_num, 2 * base):
            is_prime[i] = False

    def __init__(self, max_num):
        self.cnt = 0
        self.max_num = max_num
        max_num += 1
        max_num_root = int(math.sqrt(max_num))
        is_prime = [False] * (max_num + 30)
        self.prime = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]       # Primes smaller than 30

        for i in range(30, max_num, 30):                        # Number that is likely to be a prime
            for j in [1, 7, 11, 13, 17, 19, 23, 29]:            # Modulo-30 remainders that primes have 
                is_prime[j + i] = True

        for i in self.prime:
            self._sieve(i, max_num, is_prime)

        for i in range(30, max_num_root, 30):                   # Sieve of Eratosthenes
            for x in [1, 7, 11, 13, 17, 19, 23, 29]:
                x = x + i
                if is_prime[x]:
                    self._sieve(x, max_num, is_prime)

        for i in range(30, max_num):
            if is_prime[i]: self.prime.append(i)

    def __iter__(self):
        return self

    def __next__(self):
        if self.cnt == len(self.prime) or self.prime[self.cnt] > self.max_num:
            raise StopIteration
        else:
            self.cnt += 1
            return self.prime[self.cnt - 1]

def qpow(b, e, p):
    """
    Calculate b^e mod p
    """
    ans = 1
    while e:
        if e & 1: ans = ans * b % p
        b = b * b % p
        e >>= 1
    return ans

def miller_rabin(p, a, k, s):
    """
    Miller-Rabin prime test, use a to test p. 
    """
    x = p - 1
    b = qpow(a, s, p)
    if b == 1 or b == x: return True
    for t in range(1, k):
        b = b * b % p
        if b == x: return True
    return False

def prime_seq(start = 2):
    """
    Return a prime sequence begins from start.
    """
    for i in [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]:             # Primes less than 30
        if i < start: continue
        yield i

    if start < 30: base = 0
    else: base = (start // 30 - 1) * 30                         # Find a base so that start = 30*(base+1) + r(0<=r<30)
    while True:
        base += 30
        for i in [1, 7, 11, 13, 17, 19, 23, 29]:                # Modulo-30 remainders that primes have to create a possible prime greater than 30
            p = base + i
            if p < start: continue

            s = p - 1
            k = 0
            while not (s & 1):          # Make p-1 = 2^k*s
                k += 1
                s >>= 1
            is_prime = True

            for a in [2, 3, 5, 7, 11, 325, 9375, 28178, 450775, 9780504, 19760817, 2147483647, 17952650222]:  # Bases used for Miller-Rabin test at least suitable for number less than 2^64
                if a >= p: break
                if not miller_rabin(p, a, k, s):
                    is_prime = False
                    break
            if is_prime: yield p

print("Small Prime Test:")
error = False
primes = prime_seq()
for i in PrimeSeq(100):
    j = next(primes)
    if j != i:
        error = True
        print(f"Not match! {i} != {j}")
    else: print('%-4s' %(i), end="")
if error: print("Error!")
print("\n")

print("Large Prime Test:")
for i in islice(prime_seq(1976081719760817), 0, 5):
    print('%-20s' %(i), end="")