Python: Get bits from an list of integers and fill a bytearray

Question

I am trying in python to get the last 2 bits of a list of integers to a bytearray, by calling specific external libraries (namely sys + math) & get the bytearray printed. Basically I'm trying to use python for Elias-Fano compression method.

The problem is that I can't get them right in the bytearray... So far I got this:

import sys
import math

OriginalList = []
with open(sys.argv[1], 'r') as file:
    OriginalList = [line.rstrip() for line in file]


maxList = int(max(OriginalList))
numList = len(OriginalList)
l = math.floor(math.log2(maxList/numList))
print ("L = " + str(l))


L = bytearray()
for x in OriginalList:
  a = bin(x)[(-1-l):-1]
  L.append(a)


print ("L array")
for x in L_array:
  print(x)

Answer 1

I'm not sure why you'd want to use a bytearray for this, since you need to concatenate bit fields of different sizes. Here's a solution that just uses shifts and masks to do the bit manipulation, and stores the result in a single variable. Remember that Python integers can be of arbitrary size.

You can see that the first loop builds up the first section of the encoding, where we encode the number of elements that have the prefix. The second loop adds the fixed size leftovers.

import sys
import math

#OriginalList = []
#with open(sys.argv[1], 'r') as file:
#    OriginalList = [int(line.rstrip()) for line in file]

OriginalList = (2,3,5,7,11,13,24)

maxList = max(OriginalList)
numList = len(OriginalList)

m = math.ceil(math.log2(maxList))
n = math.ceil(math.log2(numList))
l = m-n

mask_n = (1<<n)-1
mask_l = (1<<l)-1

mycopy = list(OriginalList)
bits = 0
for i in range(1<<n):
    while mycopy and ((mycopy[0]>>l) & mask_n) == i:
        bits = (bits << 1) + 1
        mycopy.pop(0)
    bits <<= 1
for i in OriginalList:
    bits = (bits << l) | (i & mask_l)

print(bin(bits))