Add own bitarray and Miller-Rabin primality test.
This commit is contained in:
39
python/lib_bitarray.py
Normal file
39
python/lib_bitarray.py
Normal file
@@ -0,0 +1,39 @@
|
|||||||
|
class Bitarray:
|
||||||
|
def __init__(self, num_bits):
|
||||||
|
self.word_size = 32
|
||||||
|
self.word_mask = 2**self.word_size - 1
|
||||||
|
num_words = num_bits // self.word_size
|
||||||
|
if num_bits % self.word_size != 0:
|
||||||
|
num_words += 1
|
||||||
|
self.num_bits = num_bits
|
||||||
|
self.num_words = num_words
|
||||||
|
self.array = [0] * num_words
|
||||||
|
|
||||||
|
def setall(self, target: bool):
|
||||||
|
if target:
|
||||||
|
self.array = [self.word_mask] * self.num_words
|
||||||
|
else:
|
||||||
|
self.array = [0] * self.num_words
|
||||||
|
|
||||||
|
def set(self, bit, target: bool):
|
||||||
|
assert(bit < self.num_bits)
|
||||||
|
word_index = bit // self.word_size
|
||||||
|
bit_index = bit % self.word_size
|
||||||
|
value = self.array[word_index]
|
||||||
|
if target:
|
||||||
|
value = (value | (1 << bit_index))
|
||||||
|
else:
|
||||||
|
value = (value & (~(1 << bit_index)))
|
||||||
|
self.array[word_index] = value
|
||||||
|
|
||||||
|
def get(self, bit) -> bool:
|
||||||
|
assert(bit < self.num_bits)
|
||||||
|
word_index = bit // self.word_size
|
||||||
|
bit_index = bit % self.word_size
|
||||||
|
return bool(self.array[word_index] & (1 << bit_index))
|
||||||
|
|
||||||
|
def __getitem__(self, bit: int) -> bool:
|
||||||
|
return self.get(bit)
|
||||||
|
|
||||||
|
def __setitem__(self, bit: int, value: bool):
|
||||||
|
self.set(bit, value)
|
||||||
@@ -1,4 +1,5 @@
|
|||||||
from functools import lru_cache
|
from functools import lru_cache
|
||||||
|
from lib_bitarray import Bitarray
|
||||||
try:
|
try:
|
||||||
from lib_misc import get_item_counts
|
from lib_misc import get_item_counts
|
||||||
from lib_misc import product
|
from lib_misc import product
|
||||||
@@ -65,6 +66,29 @@ def is_prime(n):
|
|||||||
return True
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
def is_prime_rabin_miller(number):
|
||||||
|
""" Rabin-Miller Primality Test """
|
||||||
|
if number % 6 not in [1,5]:
|
||||||
|
return False
|
||||||
|
|
||||||
|
r, s = 1, number - 1
|
||||||
|
while s % 2 == 0:
|
||||||
|
s //= 2
|
||||||
|
r += 1
|
||||||
|
|
||||||
|
for witness in [2, 3, 5, 7, 11, 13, 17, 23, 29, 31, 37, 41, 43, 47, 53]:
|
||||||
|
remainder = pow(witness, s, number)
|
||||||
|
if remainder == 1:
|
||||||
|
continue
|
||||||
|
for _ in range(1, r):
|
||||||
|
if remainder == number - 1:
|
||||||
|
break
|
||||||
|
remainder = pow(remainder, 2, number)
|
||||||
|
else:
|
||||||
|
return False
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
def prime_nth(n):
|
def prime_nth(n):
|
||||||
"""Returns the nth prime number. The first number
|
"""Returns the nth prime number. The first number
|
||||||
is 1 indexed, i.e. n = 1 -> 2, n = 2 -> 3, etc.
|
is 1 indexed, i.e. n = 1 -> 2, n = 2 -> 3, etc.
|
||||||
@@ -97,12 +121,8 @@ def primes(n_max):
|
|||||||
|
|
||||||
:param n_max:
|
:param n_max:
|
||||||
"""
|
"""
|
||||||
try:
|
b = Bitarray(n_max)
|
||||||
import bitarray
|
|
||||||
b = bitarray.bitarray(n_max)
|
|
||||||
b.setall(True)
|
b.setall(True)
|
||||||
except ModuleNotFoundError:
|
|
||||||
b = [True for _ in range(n_max)]
|
|
||||||
n = 1
|
n = 1
|
||||||
b[n - 1] = False
|
b[n - 1] = False
|
||||||
while n * n <= n_max:
|
while n * n <= n_max:
|
||||||
|
|||||||
Reference in New Issue
Block a user