在 NumPy 数组操作中避免双循环 [英] Avoiding double for-loops in NumPy array operations
问题描述
假设我有两个二维 NumPy 数组 A
和 B
,我想计算矩阵 C
,其条目为 C[i, j] = f(A[i, :], B[:, j])
,其中 f
是一些函数,它接受两个一维数组并返回一个数字.
Suppose I have two 2D NumPy arrays A
and B
, I would like to compute the matrix C
whose entries are C[i, j] = f(A[i, :], B[:, j])
, where f
is some function that takes two 1D arrays and returns a number.
例如,如果 def f(x, y): return np.sum(x * y)
那么我只会有 C = np.dot(A, B)代码>.但是,对于一般函数
f
,是否有我可以利用的 NumPy/SciPy 实用程序比执行双 for 循环更有效?
For instance, if def f(x, y): return np.sum(x * y)
then I would simply have C = np.dot(A, B)
. However, for a general function f
, are there NumPy/SciPy utilities I could exploit that are more efficient than doing a double for-loop?
例如取def f(x, y): return np.sum(x != y)/len(x)
,其中x
和y
不是简单的 0/1 位向量.
For example, take def f(x, y): return np.sum(x != y) / len(x)
, where x
and y
are not simply 0/1-bit vectors.
推荐答案
以下是使用广播的合理通用方法.
Here is a reasonably general approach using broadcasting.
首先,将两个矩阵重塑为四阶张量.
First, reshape your two matrices to be rank-four tensors.
A = A.reshape(A.shape + (1, 1))
B = B.reshape((1, 1) + B.shape)
其次,一个元素一个元素地应用你的函数,而不执行任何归约.
Second, apply your function element by element without performing any reduction.
C = f(A, B) # e.g. A != B
重构矩阵后,numpy 可以广播.结果张量 C
具有形状 A.shape + B.shape
.
Having reshaped your matrices allows numpy to broadcast. The resulting tensor C
has shape A.shape + B.shape
.
第三,应用任何所需的减少,例如,对要丢弃的索引求和:
Third, apply any desired reduction by, for example, summing over the indices you want to discard:
C = C.sum(axis=(1, 3)) / C.shape[0]
这篇关于在 NumPy 数组操作中避免双循环的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持IT屋!