40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
|
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)
|