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