1from bitarray import bitarray
2from bitarray.util import int2ba, ba2int, pprint
3
4
5class SmallIntArray(object):
6    """
7    A class which allows efficiently storeing an array of integers
8    represented by a specified number of bits.
9    For example, an array with 1000 5 bit integers can be created,
10    allowing each element in the array to take values form 0 to 31,
11    while the size of the object is 625 (5000/8) bytes.
12    """
13    def __init__(self, N, k):
14        self.N = N  # number of integers
15        self.k = k  # bits for each integer
16        self.array = bitarray(N * k)
17
18    def slice_i(self, i):
19        assert 0 <= i < self.N
20        return slice(self.k * i, self.k * (i + 1))
21
22    def __getitem__(self, i):
23        return ba2int(self.array[self.slice_i(i)])
24
25    def __setitem__(self, i, v):
26        self.array[self.slice_i(i)] = int2ba(v, self.k)
27
28
29if __name__ == '__main__':
30    from random import randint
31
32    # define array with 1000 integers, each represented by 5 bits
33    a = SmallIntArray(1000, 5)
34
35    b = [] # store values, for assertion below
36    for i in range(1000):
37        v = randint(0, 31)
38        b.append(v)
39        a[i] = v
40
41    print(b[:5])
42    print(a.array.buffer_info())
43    pprint(a.array[:25], group=5)
44
45    for i in range(1000):
46        assert a[i] == b[i]
47