优化整数向量转二进制向量的处理性能

阿华AIGC实验室

2026-4-14

优化整数向量转二进制向量的处理性能

嘿，我完全懂你现在的头疼——400万条向量，每条254个元素，用Python循环加format处理确实会慢到让人崩溃。原代码里的两层Python循环（外层遍历数据、内层遍历每个向量元素）加上纯Python的map和format操作，在大数据量下就是性能杀手。咱们直接用numpy的向量化操作来重构，这能带来数量级的性能提升，同时还能减少内存占用。

原代码的性能瓶颈分析

原代码的核心问题在于：

依赖Python级别的循环，而Python循环的执行效率远低于numpy底层的C实现；
format和map都是纯Python操作，无法利用numpy的向量化优化；
每次循环都要创建新的列表和数组，产生大量不必要的内存开销。

优化方案1：使用numpy `unpackbits`（最快最简洁）

因为你的整数范围是0-65535（正好是16位无符号整数uint16的范围），我们可以直接利用numpy的unpackbits函数——它专门用于将字节拆分成二进制位，是底层优化的极致：

import numpy as np

def integer_vectors_to_binary_fast(data, bits=16):
    # 将输入转换为2D uint16数组（如果输入是列表的话）
    if not isinstance(data, np.ndarray):
        data = np.array(data, dtype=np.uint16)
    # 将uint16元素拆分为两个uint8字节（大端序，保证高位在前）
    byte_arr = data.view(np.uint8).reshape(data.shape + (2,))
    # 拆分每个字节为8位二进制，然后扁平化结果
    binary_arr = np.unpackbits(byte_arr, axis=-1).reshape(data.shape[0], -1)
    return binary_arr

优化方案2：向量化位运算（更灵活，支持任意位数）

如果你以后需要处理非16位的情况，或者需要更灵活的位操作，可以用掩码+位运算的方式，同样是完全向量化的：

import numpy as np

def integer_vectors_to_binary_bitwise(data, bits=16):
    if not isinstance(data, np.ndarray):
        data = np.array(data, dtype=np.uint16)
    # 生成从最高位到最低位的掩码数组（比如16位就是[32768, 16384, ..., 1]）
    masks = 2 ** np.arange(bits-1, -1, -1, dtype=np.uint16)
    # 对所有元素同时进行位运算，提取每一位并转换为int
    binary_arr = ((data[..., None] & masks) != 0).astype(np.uint8).reshape(data.shape[0], -1)
    return binary_arr

测试验证

用你给出的示例测试，两个优化函数的输出和原代码完全一致：

vec_a = np.asarray([12, 15, 14])
print(integer_vectors_to_binary_fast([vec_a]))
# 输出：array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0]], dtype=uint8)