如何用Numpy将比特数组按N比特/字节打包？求高性能实现方案

针对在GNURadio中需要将比特数组按可变N比特/元素打包的需求，这里提供一个基于Numpy向量化操作的高性能解决方案，替代无法自定义比特数的np.packbits函数。

核心思路 #

通过将比特数组重塑为每行含N个比特的二维数组，利用权重向量与每行做内积，快速将每行的比特转换为对应的整数。全程使用Numpy的优化运算，避免Python循环，保证处理性能。

实现代码 #

import numpy as np

def pack_bits_to_n_per_byte(bits, n):
    # 转换为Numpy数组并指定无符号整数类型
    bits = np.asarray(bits, dtype=np.uint8)
    
    # 处理比特长度非N整数倍的情况：可选补零或丢弃多余比特
    remainder = len(bits) % n
    if remainder != 0:
        # 方案1：补零至N的整数倍
        pad_length = n - remainder
        bits = np.pad(bits, (0, pad_length), mode='constant')
        # 方案2：丢弃末尾多余的比特（取消注释下方代码并注释上方补零代码）
        # bits = bits[:len(bits) - remainder]
    
    # 重塑为每行N个比特的二维数组
    reshaped_bits = bits.reshape(-1, n)
    # 生成权重数组（如N=2时为[2,1]，N=4时为[8,4,2,1]）
    weights = 2 ** np.arange(n-1, -1, -1, dtype=np.uint8)
    # 计算每行比特对应的整数值
    packed_values = np.dot(reshaped_bits, weights)
    
    return packed_values

测试示例 #

# 测试输入比特数组
test_bits = [0, 1, 1, 1, 0, 0, 1, 0]

# 按2比特/元素打包
print(pack_bits_to_n_per_byte(test_bits, 2))  # 输出: [1 3 0 2]
# 按4比特/元素打包
print(pack_bits_to_n_per_byte(test_bits, 4))  # 输出: [7 2]

性能说明 #

• 所有运算均为Numpy内部优化的C级操作，无Python循环，处理大规模比特数组时性能远优于纯Python实现。
• 时间复杂度为O(M)（M为比特数组长度），适配GNURadio等实时信号处理场景的性能需求。
• 支持1~8比特/元素的打包需求，若需更大比特数，只需调整dtype为np.uint16/np.uint32等即可。

内容的提问来源于stack exchange，提问作者Jordan