Python, быстрое сжатие большого количества чисел с Elias Gamma

Question

У нас есть 2d список, при необходимости мы можем преобразовать его во что угодно. Каждая строка содержит некоторые положительные целые числа (дельты исходных возрастающих чисел). Всего 2 миллиарда чисел, более половины из которых равно 1. При использовании кодирования Elias-Gamma мы можем кодировать 2d список построчно (позже мы будем обращаться к произвольным строкам с индексом строки), используя около 3 бита на число на основе расчета из распределения. Однако наша программа работает уже 12 часов и до сих пор не завершила кодирование. Вот что мы делаем:

from bitstring import BitArray
def _compress_2d_list(input: List[List[int]]) -> List[BitArray]:
    res = []
    for row in input:
        res.append(sum(_elias_gamma_compress_number(num) for num in row))
    return res 

def _elias_gamma_compress_number(x: int) -> BitArray:
    n = _log_floor(x)
    return BitArray(bin="0" * n) + BitArray(uint=x, length=_log_floor(x) + 1)

def log_floor(num: int) -> int:
    return floor(log(num, 2))

Вызывается:

input_2d_list: List[List[int]]  # containing 1.5M lists, total 2B numbers
compressed_list = _compress_2d_list(input_2d_list)

Как я могу оптимизировать свой код, чтобы он работал быстрее? Я имею в виду, НАМНОГО БЫСТРЕЕ ... Я согласен с использованием любой надежной популярной библиотеки или структуры данных.

Кроме того, как нам быстрее распаковать с помощью BitStream? В настоящее время я читаю префикс 0 один за другим, а затем читаю двоичный файл сжатого числа через некоторое время l oop. Тоже не очень быстро ...

Answer 1

Если вас устраивают numpy "битовые поля", вы можете выполнить сжатие за считанные минуты. Декодирование происходит в три раза медленнее, но все же это вопрос минут.

Пример выполнения:

# create example (1'000'000 numbers) 
a = make_example()
a
# array([2, 1, 1, ..., 3, 4, 3])

b,n = encode(a) # takes ~100 ms on my machine
c = decode(b,n) #       ~300 ms

# check round trip
(a==c).all()
# True

Код:

import numpy as np
    
def make_example():
    a = np.random.choice(2000000,replace=False,size=1000001)
    a.sort()
    return np.diff(a)

def encode(a):
    a = a.view(f'u{a.itemsize}')
    l = np.log2(a).astype('u1')
    L = ((l<<1)+1).cumsum()
    out = np.zeros(L[-1],'u1')
    for i in range(l.max()+1):
        out[L-i-1] += (a>>i)&1
    return np.packbits(out),out.size

def decode(b,n):
    b = np.unpackbits(b,count=n).view(bool)
    s = b.nonzero()[0]
    s = (s<<1).repeat(np.diff(s,prepend=-1))
    s -= np.arange(-1,len(s)-1)
    s = s.tolist() # list has faster __getitem__
    ns = len(s)
    def gen():
        idx = 0
        yield idx
        while idx < ns:
            idx = s[idx]
            yield idx
    offs = np.fromiter(gen(),int)
    sz = np.diff(offs)>>1
    mx = sz.max()+1
    out = np.zeros(offs.size-1,int)
    for i in range(mx):
        out[b[offs[1:]-i-1] & (sz>=i)] += 1<<i
    return out

Answer 2

Некоторые простые оптимизации приводят к улучшению в три раза:

def _compress_2d_list(input):
    res = []
    for row in input:
        res.append(BitArray('').join(BitArray(uint=x, length=2*x.bit_length()-1) for x in row))
    return res

Однако я думаю, вам понадобится что-то получше. На моей машине это будет sh примерно за 12 часов на 1,5 миллионах списков с 1400 дельтами в каждом.

В C кодирование занимает около минуты. Около 15 секунд на декодирование.