151 lines
3.3 KiB
Python
151 lines
3.3 KiB
Python
from functools import lru_cache
|
|
try:
|
|
from lib_misc import get_item_counts
|
|
from lib_misc import product
|
|
except ModuleNotFoundError:
|
|
from .lib_misc import get_item_counts
|
|
from .lib_misc import product
|
|
|
|
|
|
def prime_factors(n):
|
|
"""
|
|
Returns a list of prime factors for n.
|
|
|
|
:param n: number for which prime factors should be returned
|
|
"""
|
|
# TODO: Look into using a prime wheel instead.
|
|
factors = []
|
|
rest = n
|
|
divisor = 2
|
|
while rest % divisor == 0:
|
|
factors.append(divisor)
|
|
rest //= divisor
|
|
divisor = 3
|
|
while divisor * divisor <= rest:
|
|
while rest % divisor == 0:
|
|
factors.append(divisor)
|
|
rest //= divisor
|
|
divisor += 2
|
|
if rest != 1:
|
|
factors.append(rest)
|
|
return factors
|
|
|
|
|
|
def prime_factors_count(n):
|
|
"""
|
|
Returns a dictionay of primes where each key is a prime
|
|
and the value how many times that prime is part of the factor of n.
|
|
|
|
:param n: numober for which prime factor counts are returned
|
|
:returns: a dict where they key is a prime and the value
|
|
the count of how often that value occurs
|
|
|
|
"""
|
|
return get_item_counts(prime_factors(n))
|
|
|
|
|
|
@lru_cache(maxsize=10000)
|
|
def is_prime(n):
|
|
"""Returns True if n is prime and False otherwise.
|
|
|
|
:param n: number to be checked
|
|
|
|
"""
|
|
if n < 2:
|
|
return False
|
|
if n == 2 or n == 3:
|
|
return True
|
|
if n % 2 == 0:
|
|
return False
|
|
d = 3
|
|
while d * d <= n:
|
|
if n % d == 0:
|
|
return False
|
|
d += 2
|
|
return True
|
|
|
|
|
|
def prime_nth(n):
|
|
"""Returns the nth prime number. The first number
|
|
is 1 indexed, i.e. n = 1 -> 2, n = 2 -> 3, etc.
|
|
|
|
:param n:
|
|
|
|
"""
|
|
if n == 1:
|
|
return 2
|
|
if n == 2:
|
|
return 3
|
|
p = 5
|
|
i = 3
|
|
while i < n:
|
|
p = p + 2
|
|
if is_prime(p):
|
|
i += 1
|
|
return p
|
|
|
|
|
|
def primes(n_max):
|
|
"""
|
|
Returns a list of all primes smaller n based
|
|
on sieve of erasthones algorithm.
|
|
|
|
If bitarray module is installed it will be used
|
|
for the array, otherwise a regular Python list is
|
|
used. The function should work for much larger
|
|
numbers with bitarray.
|
|
|
|
:param n_max:
|
|
"""
|
|
try:
|
|
import bitarray
|
|
b = bitarray.bitarray(n_max)
|
|
b.setall(True)
|
|
except ModuleNotFoundError:
|
|
b = [True for _ in range(n_max)]
|
|
n = 1
|
|
b[n - 1] = False
|
|
while n * n <= n_max:
|
|
if b[n - 1] is True:
|
|
for i in range(n + n, n_max + 1, n):
|
|
b[i - 1] = False
|
|
n += 1
|
|
ps = []
|
|
for i in range(1, n_max + 1):
|
|
if b[i - 1]:
|
|
ps.append(i)
|
|
return ps
|
|
|
|
|
|
def get_divisors_count(n):
|
|
"""
|
|
Returns the number of divisors for n.
|
|
The numbers 1 and n count as a divisor.
|
|
|
|
>>> get_divisors_count(1)
|
|
1
|
|
>>> get_divisors_count(3)
|
|
2 # 1, 3
|
|
>>> get_divisors_count(4)
|
|
3 # 1, 2, 4
|
|
|
|
Getting the number of divisors is a combinatorial
|
|
problem that can be solved by using the counts
|
|
for each prime factor. For example, consider
|
|
|
|
2 * 2 * 7 = 28
|
|
|
|
We have 3 options for 2 (1, 1 * 2, 2 * 2)
|
|
and 2 options for 7 (1, 1 * 7).
|
|
|
|
By multiplying those options we get the number
|
|
of combinations:
|
|
|
|
2 * 3 = 6
|
|
"""
|
|
if n == 1:
|
|
return 1
|
|
factors = prime_factors_count(n)
|
|
count = product([v + 1 for v in factors.values()])
|
|
return count
|