Numpy布尔二维数组与float64向量点积内存过高问题及优化咨询

问题原因 #

Numpy的np.dot依赖BLAS/LAPACK库实现高效运算，而这些库仅针对浮点类型做了深度优化。当你传入布尔二维数组a和float64一维数组v时，Numpy会先把整个布尔数组转换成float64类型，再执行点积。转换后的数组大小为286880 * 20419 * 8字节 ≈ 43.6GiB，远超系统内存上限，直接触发内存错误。

虽然理论上点积只需逐行累加对应元素，但Numpy的通用实现优先追求运算速度，没有针对布尔数组“仅需累加True对应位置值”的特性做特殊优化。

优化解决方案 #

1. 稀疏矩阵法（内存最优） #

布尔数组的核心特点是大量元素为False，转成稀疏矩阵后仅存储True元素的位置，内存占用骤降：

import numpy as np
from scipy.sparse import csr_matrix

x, y = 286880, 20419
a = np.random.randint(0, 2, (x, y), dtype=np.bool_)
v = np.random.rand(y)

# 转换为CSR格式稀疏矩阵，自动忽略False元素
a_sparse = csr_matrix(a)
# 执行点积运算
d = a_sparse.dot(v)

CSR矩阵仅存储非零元素的索引和值，内存占用仅与True元素数量相关，远低于原布尔数组的大小。

2. Numba加速逐行累加（时间与内存平衡） #

用Numba编译循环代码，直接利用布尔数组特性逐行累加，无需转换整个数组，同时保持接近原生的运算速度：

import numpy as np
from numba import njit

x, y = 286880, 20419
a = np.random.randint(0, 2, (x, y), dtype=np.bool_)
v = np.random.rand(y)

@njit(parallel=True)  # 启用多线程并行，利用多核加速
def bool_matrix_dot(a, v):
    x, y = a.shape
    result = np.zeros(x, dtype=np.float64)
    for i in range(x):
        row_sum = 0.0
        for j in range(y):
            if a[i, j]:
                row_sum += v[j]
        result[i] = row_sum
    return result

d = bool_matrix_dot(a, v)

这种方式仅需存储原布尔数组、一维数组和结果数组，总内存占用不到2GB，同时通过多线程加速弥补循环的性能损失。

3. 纯Numpy列累加（无依赖） #

如果不想用第三方库，可以用纯Numpy的列遍历累加，内存占用极低，速度略慢但足够实用：

import numpy as np

x, y = 286880, 20419
a = np.random.randint(0, 2, (x, y), dtype=np.bool_)
v = np.random.rand(y)

d = np.zeros(x, dtype=np.float64)
for j in range(y):
    # 将v[j]加到所有a[:,j]为True的行
    d[a[:, j]] += v[j]

每次循环仅处理一列的掩码，无需转换整个布尔数组，内存压力极小。

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