如何用Numpy实现自定义2D稀疏二进制矩阵运算替代慢速循环

核心实现思路 #

你需要的运算可以直接通过numpy的广播机制完成，无需任何循环：将形状为(I,J)的矩阵A扩展一个最后维度变为(I,J,1)，将形状为(J,K)的矩阵B扩展一个最前维度变为(1,J,K)，两者逐元素相乘/逻辑与后会自动广播为(I,J,K)的结果，刚好满足「仅当A(i,j)=1且B(j,k)=1时C(i,j,k)=1」的规则。

代码实现 #

import numpy as np

I = 2
J = 3
K = 4

A = np.random.randint(2, size=(I, J))
B = np.random.randint(2, size=(J, K))

# 广播乘法实现，一行代替三层循环
C = A[:, :, np.newaxis] * B[np.newaxis, :, :]

# 也可以用逻辑与实现，性能更高
# C = np.logical_and(A[:, :, None], B[None, :, :]).astype(np.uint8)

# 验证和三层循环结果完全一致
C_loop = np.zeros((I,J,K))
for i in range(I):
    for j in range(J):
        for k in range(K):
            if A[i,j] == 1 and B[j,k] == 1:
                C_loop[i,j,k] = 1
print(np.allclose(C, C_loop)) # 输出True

关于np.multiply.outer调用失败的说明 #

你之前直接调用np.multiply.outer(A,B)得到的张量形状为(I,J,J,K)，会多出来一个冗余的J维度，只有第二个和第三个J维度对齐的位置才是你需要的结果，你可以通过如下写法调整得到正确结果，但可读性不如广播实现：

C = np.multiply.outer(A,B).diagonal(axis1=1, axis2=2).transpose(0,2,1)

性能对比 #

当I/J/K数值较大时，广播实现的性能优势极为明显：例如I=200、J=300、K=400的场景下，三层循环实现需要数秒才能跑完，而广播实现仅需几毫秒，速度提升可达数千倍。

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