imp/celeste/math/primes/__init__.py

48 lines
1.5 KiB
Python

from math import inf, isqrt # integer square root
from itertools import takewhile, compress
SMALL_PRIMES = (2,3,5,7,11,13,17,19,23,29,31,37,41,43,47,53,59)
'''
Euler's Totient (Phi) Function
Implemented in O(nloglog(n)) using the Sieve of Eratosthenes.
'''
def eulertotient(n: int) -> int:
phi = int(n > 1 and n)
for p in range(2, isqrt(n) + 1):
if not n % p:
phi -= phi // p
while not n % p:
n //= p
#if n is > 1 it means it is prime
if n > 1: phi -= phi // n
return phi
'''
Tests the primality of an integer using its totient.
NOTE: If totient(n) has already been calculated
then pass it as the optional phi parameter.
'''
def is_prime(n: int, phi: int = None) -> bool:
return n - 1 == (phi if phi is not None else eulertotient(n))
# Taken from Lucas A. Brown's primefac.py (some variables renamed)
def primegen(limit=inf) -> int:
ltlim = lambda x: x < limit
yield from takewhile(ltlim, SMALL_PRIMES)
pl, prime = [3,5,7], primegen()
for p in pl: next(prime)
n = next(prime); nn = n*n
while True:
n = next(prime); ll, nn = nn, n*n
delta = nn - ll
sieve = bytearray([True]) * delta
for p in pl:
k = (-ll) % p
sieve[k::p] = bytearray([False]) * ((delta-k)//p + 1)
if nn > limit: break
yield from compress(range(ll,ll+delta,2), sieve[::2])
pl.append(n)
yield from takewhile(ltlim, compress(range(ll,ll+delta,2), sieve[::2]))