避免在NumPy数组操作中使用双for循环

Question

假设我有两个二维NumPy数组A和B ，我想计算矩阵C其条目为C[i, j] = f(A[i, :], B[:, j]) ，其中f是该函数需要两个一维数组并返回一个数字。

例如，如果def f(x, y): return np.sum(x * y)那么我将简单地拥有C = np.dot(A, B) 。 但是，对于通用函数f ，我可以利用NumPy / SciPy实用程序比执行双for循环更有效吗？

例如，取def f(x, y): return np.sum(x != y) / len(x) ，其中x和y不仅仅是0/1位向量。

Answer 1

这是使用广播的一种合理的通用方法。

首先，将两个矩阵重塑为四阶张量。

A = A.reshape(A.shape + (1, 1))
B = B.reshape((1, 1) + B.shape)

其次，逐个应用函数，而不执行任何缩减。

C = f(A, B)  # e.g. A != B

重塑矩阵后，numpy可以广播 。 所得张量C形状为A.shape + B.shape 。

第三，例如通过对要丢弃的索引求和来应用任何所需的减少量：

C = C.sum(axis=(1, 3)) / C.shape[0]