如何用numpy向量化实现基于指示矩阵C的A行与B列加减操作？

当然有更高效的向量化实现方式！完全不用写嵌套循环，利用NumPy的矩阵乘法就能一步完成，而且速度会比循环快得多——毕竟NumPy的底层是用优化过的C代码实现的，尤其适合这种大规模矩阵操作。

核心思路 #

你的需求本质上是根据指示矩阵C的权重，把A的行线性组合后加到B的对应列上。我们可以用矩阵乘法直接完成这个操作：
最终更新后的B可以表示为：B = B + A.T @ C

这里的关键是形状匹配：

• A是(n,m)，转置后A.T是(m,n)
• C是(n,k)
• 两者相乘得到(m,k)的矩阵，正好和B的形状(m,k)完全一致，可以直接和B做加法。

示例代码 #

下面用具体例子验证这个方法，顺便和循环版本对比结果：

import numpy as np

# 构造测试用的矩阵
n, m, k = 3, 4, 5
A = np.arange(n*m).reshape(n, m)  # 形状(3,4)的输入矩阵
B = np.ones((m, k))               # 初始的(4,5)矩阵B
C = np.array([[-1, 0, 0, 0, 1],
              [0, 0, 0, 0, -1],
              [0, 1, -1, 0, 0]])  # 指示矩阵，形状(3,5)

# 向量化实现（一行搞定）
B_updated = B + A.T @ C

# 循环实现（用于验证结果）
B_loop = B.copy()
for i in range(n):
    for j in range(k):
        if C[i,j] == 1:
            B_loop[:, j] += A[i, :]
        elif C[i,j] == -1:
            B_loop[:, j] -= A[i, :]

# 检查两种方法结果是否一致
print(np.allclose(B_updated, B_loop))  # 输出 True，说明完全一致

为什么这个方法可行？ #

矩阵乘法A.T @ C会自动帮我们完成所有的加权求和：
对于结果矩阵中的每个元素(A.T @ C)[p, j]，它等于sum(A[i,p] * C[i,j] for i in range(n))——这正好对应：

• 当C[i,j] = 1时，把A第i行的第p个元素加到B第p行第j列
• 当C[i,j] = -1时，从B第p行第j列减去A第i行的第p个元素
• 当C[i,j] = 0时，这个元素对结果没有贡献

完全符合你的需求，而且全程没有Python循环，效率拉满！

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