Python，使用 Elias Gamma 快速压缩大量数字

Question

我们有一个二维列表，如有必要，我们可以将其转换为任何内容。每行包含一些正整数（原始递增数字的增量）。总共 20 亿个数字，其中一半以上等于 1。当使用Elias-Gamma 编码时，我们可以根据每个数字使用大约 3 位逐行编码 2d 列表（稍后我们将使用行索引访问任意行）从分布计算。然而，我们的程序已经运行了 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，然后在while循环中读取压缩数字的二进制文件。也不是很快……

Answer 1

如果你对numpy“bitfields”没问题，你可以在几分钟内完成压缩。解码速度要慢三倍，但仍然需要几分钟。

示例运行：

# 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

但是，我认为您需要比这更好的东西。在我的机器上，这将在大约 12 小时内完成 150 万个列表，每个列表有 1400 个增量。

在 C 语言中，编码大约需要一分钟。解码大约需要 15 秒。