Python自定义Bitarray类的高效位运算实现优化咨询

问题背景 #

我自行实现了一个处理位值数组的Bitarray类，当前通过遍历布尔值列表执行位运算，效率极低。需求如下：

• 基于随机密钥和XOR运算构建加密协议，需自行实现（不依赖cryptography、bitarray等第三方库）
• 支持单个位值的访问
• 无需频繁转换为bytes即可执行位运算

原实现代码如下：

import sys

class Bitarray:
    """ Representation of an array of bits.
    :param bits: the list of boolean values (i.e. {False, True}) of the bitarray.
    """
    def __init__(self, values:list[bool]):
        self._bits:list[bool] = values
        self._length:int = len(values)

    @staticmethod
    def from_bytes(data:bytes, byteorder:str = None):
        def _access_bit(data, index):
            """ Credits: https://stackoverflow.com/a/43787831/23022499 """
            base = int(index // 8)
            shift = int(index % 8)
            return (data[base] >> shift) & 0x1
        
        if byteorder == None:
            byteorder = sys.byteorder
        elif byteorder != 'little' and byteorder != 'big':
            raise ValueError('Param byteorder must be either "little" or "big".')
        
        bin_data = [_access_bit(data, i) for i in range(len(data) * 8)]
        bin_data = [bool(b) for b in bin_data]
        return Bitarray(bin_data) if byteorder == 'big' else Bitarray(bin_data[::-1])
    
    def __getitem__(self, index) -> bool:
        return self._bits[index]
    
    def __len__(self) -> int:
        return self._length
    
    # bit-wise operations
    def __and__(self, other):
        if type(other) != Bitarray:
            raise TypeError("Unsupported operand type(s) for &: '{}' and '{}'".format(type(self), type(other)))
        
        if self._length != len(other):
            raise IndexError("The arguments for bitwise operations must have same length.")
        
        return Bitarray([(a & b) for a, b in zip(self._bits, other._bits)])

    def __or__(self, other):
        if type(other) != Bitarray:
            raise TypeError("Unsupported operand type(s) for |: '{}' and '{}'".format(type(self), type(other)))
        
        if self._length != len(other):
            raise IndexError("The arguments for bitwise operations must have same length.")
        
        return Bitarray([(a | b) for a, b in zip(self._bits, other._bits)])

    def __xor__(self, other):
        if type(other) != Bitarray:
            raise TypeError("Unsupported operand type(s) for ^: '{}' and '{}'".format(type(self).__name__, type(other).__name__))
        
        if self._length != len(other):
            raise IndexError("The arguments for bitwise operations must have same length.")
        
        return Bitarray([(a ^ b) for a, b in zip(self._bits, other._bits)])

    # to string
    def __str__(self):
        return ''.join(str(int(b)) for b in self._bits)

测试代码：

import os
import sys

a = Bitarray.from_bytes(os.urandom(16 // 8), sys.byteorder)
b = Bitarray.from_bytes(os.urandom(16 // 8), sys.byteorder)

print('XOR result: {}'.format(a ^ b))

优化方案 #

核心思路是用字节数组存储位数据，利用CPU原生的位运算指令替代Python层面的布尔值遍历，大幅提升运算效率。具体实现如下：

优化后的Bitarray类 #

import sys

class Bitarray:
    """ Representation of an array of bits, optimized for bitwise operations.
    :param data: Bytes object storing the bit data
    :param length: Total number of bits in the array
    :param byteorder: Byte order used to store the bits
    """
    def __init__(self, data: bytes, length: int, byteorder: str = sys.byteorder):
        self._data = data
        self._length = length
        self._byteorder = byteorder
        # 计算补位数量（当长度不是8的倍数时）
        self._padding = (8 - (length % 8)) % 8

    @staticmethod
    def from_bytes(data: bytes, byteorder: str = None):
        if byteorder is None:
            byteorder = sys.byteorder
        elif byteorder not in ('little', 'big'):
            raise ValueError('Param byteorder must be either "little" or "big".')
        
        length = len(data) * 8
        # 大端模式直接保留字节顺序，小端模式反转字节以匹配位序
        adjusted_data = data if byteorder == 'big' else data[::-1]
        return Bitarray(adjusted_data, length, byteorder)

    def __getitem__(self, index: int) -> bool:
        if index < 0 or index >= self._length:
            raise IndexError("Bit index out of range")
        
        # 计算目标位所在的字节索引
        byte_idx = index // 8
        # 计算位在字节内的偏移（0代表最低位）
        bit_pos = index % 8
        
        # 提取目标位
        byte = self._data[byte_idx]
        return (byte >> bit_pos) & 1 == 1

    def __len__(self) -> int:
        return self._length

    def _to_int(self) -> int:
        """ 将位数组转换为整数，用于快速位运算 """
        return int.from_bytes(self._data, byteorder=self._byteorder)

    @classmethod
    def _from_int(cls, value: int, length: int, byteorder: str = sys.byteorder):
        """ 从整数创建Bitarray，保留指定长度和字节序 """
        byte_length = (length + 7) // 8
        data = value.to_bytes(byte_length, byteorder=byteorder, signed=False)
        # 处理非8倍数长度的补位，掩码清除多余位
        if length % 8 != 0:
            mask = (1 << length) - 1
            value &= mask
            data = value.to_bytes(byte_length, byteorder=byteorder, signed=False)
        # 根据字节序调整存储的字节顺序
        adjusted_data = data if byteorder == 'big' else data[::-1]
        return cls(adjusted_data, length, byteorder)

    # 位运算实现
    def __and__(self, other):
        if not isinstance(other, Bitarray):
            raise TypeError(f"Unsupported operand type(s) for &: '{type(self).__name__}' and '{type(other).__name__}'")
        if self._length != other._length:
            raise IndexError("Bit arrays must have the same length for bitwise operations")
        
        self_int = self._to_int()
        other_int = other._to_int()
        result_int = self_int & other_int
        return Bitarray._from_int(result_int, self._length, self._byteorder)

    def __or__(self, other):
        if not isinstance(other, Bitarray):
            raise TypeError(f"Unsupported operand type(s) for |: '{type(self).__name__}' and '{type(other).__name__}'")
        if self._length != other._length:
            raise IndexError("Bit arrays must have the same length for bitwise operations")
        
        self_int = self._to_int()
        other_int = other._to_int()
        result_int = self_int | other_int
        return Bitarray._from_int(result_int, self._length, self._byteorder)

    def __xor__(self, other):
        if not isinstance(other, Bitarray):
            raise TypeError(f"Unsupported operand type(s) for ^: '{type(self).__name__}' and '{type(other).__name__}'")
        if self._length != other._length:
            raise IndexError("Bit arrays must have the same length for bitwise operations")
        
        self_int = self._to_int()
        other_int = other._to_int()
        result_int = self_int ^ other_int
        return Bitarray._from_int(result_int, self._length, self._byteorder)

    def __str__(self) -> str:
        # 生成字符串用于显示，仅在需要时遍历位
        bit_str = []
        for i in range(self._length):
            bit_str.append('1' if self[i] else '0')
        return ''.join(bit_str)

优化说明 #

1. 存储优化：用bytes对象存储位数据，避免布尔列表的内存开销和遍历成本。
2. 位运算加速：将位数组转换为整数，利用Python整数的原生位运算（底层为CPU指令）执行操作，效率比遍历布尔值提升数个数量级。
3. 单访问兼容：通过计算字节索引和位偏移提取单个位，仅需少量算术运算，性能损失可忽略。
4. 字节序兼容：保留原实现的字节序支持，确保输入输出行为一致。

测试验证 #

原测试代码无需修改，直接运行即可获得相同结果，但运算效率显著提升：

import os
import sys

a = Bitarray.from_bytes(os.urandom(2), sys.byteorder)  # 16 bits = 2 bytes
b = Bitarray.from_bytes(os.urandom(2), sys.byteorder)

print('XOR result: {}'.format(a ^ b))

超大规模场景优化 #

如果处理超大规模位数组（超出Python整数高效处理范围），可改用bytearray逐字节执行位运算，避免单个整数过大的性能损耗：

def __xor__(self, other):
    if self._length != other._length:
        raise IndexError("Bit arrays must have the same length")
    result_data = bytearray()
    for b1, b2 in zip(self._data, other._data):
        result_data.append(b1 ^ b2)
    # 处理非8倍数长度，截断补位
    if self._padding != 0:
        mask = 0xFF >> self._padding
        result_data[-1] &= mask
    return Bitarray(result_data, self._length, self._byteorder)

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