euler/python/lib_bitarray.py

40 lines
1.2 KiB
Python
Raw Normal View History

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)