CodyKochmann
1/25/2017 - 1:44 PM

Count bitarrays with collections.Counter in python

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)