Count bitarrays with collections.Counter in python
# This class was built to count bitarrays with collections.Counter
# by: codykochmann
from bitarray import bitarray
from collections import Counter
class BitarrayCounter(Counter):
"""counts bitarrays"""
def __init__(self):
super(BitarrayCounter, self).__init__()
def count(self, g):
""" counts bitarrays in a generator """
#assert isinstance(g, type(i for i in range(1)))
self.update(self.__to_int__(i) for i in g)
@staticmethod
def __to_int__(i):
""" returns a storable int from a bitarray """
return int((bitarray('1')+i).to01(), 2)
@staticmethod
def __from_int__(i):
""" returns a bitarray from a stored int """
out = bitarray()
while i>0:
out.append(i%2)
i = i/2
out.reverse()
return out[1:]
def __getitem__(self, key):
""" gets the count of the given bitarray """
# self.bit is added so any zeros at the beginning are not lost
return super(BitarrayCounter, self).__getitem__(self.__to_int__(key))
def __setitem__(self, key, value):
""" sets the count for a given bitarray """
if isinstance(key, bitarray):
return super(BitarrayCounter, self).__setitem__(self.__to_int__(key), value)
else:
return super(BitarrayCounter, self).__setitem__(key, value)
def __str__(self):
""" returns the printable version of the counter """
out = {}
for stored_int in self:
b = self.__from_int__(stored_int)
out[b] = super(BitarrayCounter, self).__getitem__(stored_int)
return '{}({})'.format(self.__class__.__name__, out)