使用 numpy 对两个数组进行迭代操作

Question

我正在使用两个不同的数组(75x4(，并且我在两个数组之间应用了最短距离算法。

所以我想：

对第一个数组的一行执行

• 操作，对第二个数组的每一行执行操作，迭代以获得 75 个值
• 找到最小值，并将其存储在新数组中
• 对第一个数组的第二行重复此操作，再次迭代第二个数组的所有行的操作，并再次存储与新数组的最小差值

我将如何使用numpy来做到这一点？

本质上，我想在数组 2 的每一行上的一行数组 1 之间执行操作，找到最小值，并将其存储在新数组中。然后对数组 1 的第 2 行执行相同的操作，对数组 1 的所有 75 行执行此类推。

这是我正在使用的公式的代码。我在这里得到的只是数组 1(训练数据(和数组 2(测试数据(的每一行之间的距离。但是我正在寻找的是对数组 1 的一行执行此操作，向下迭代数组 2 的所有行，将最小值存储在新数组中，然后对数组 1 的下一行执行相同的操作，依此类推。

arr_attributedifference = (arr_trainingdata - arr_testingdata)**2
arr_distance = np.sqrt(arr_attributedifference.sum(axis=1))

Answer 1

这里有两个方法，一个使用einsum，另一个KDTree：

einsum基本上完成了我们也可以通过广播实现的目标，例如np.einsum('ik,jk', A, B)大致相当于(A[:, None, :] * B[None, :, :]).sum(axis=2).einsum 的优点是它直接进行求和，因此避免了创建 mxmxn 中间数组。

KDTree更复杂。我们必须预先投资生成树，但之后查询最近的邻居非常有效。

import numpy as np
from scipy.spatial import cKDTree as KDTree
def f_einsum(A, B):
B2AB = np.einsum('ij,ij->i', B, B) / 2 - np.einsum('ik,jk', A, B)
idx = B2AB.argmin(axis=1)
D = A-B[idx]
return np.sqrt(np.einsum('ij,ij->i', D, D)), idx
def f_KDTree(A, B):
T = KDTree(B)
return T.query(A, 1)
m, n = 75, 4
A, B = np.random.randn(2, m, n)
de, ie = f_einsum(A, B)
dt, it = f_KDTree(A, B)
assert np.all(ie == it) and np.allclose(de, dt)
from timeit import timeit
for m, n in [(75, 4), (500, 4)]:
A, B = np.random.randn(2, m, n)
print(m, n)
print('einsum:', timeit("f_einsum(A, B)", globals=globals(), number=1000))
print('KDTree:', timeit("f_KDTree(A, B)", globals=globals(), number=1000))

示例运行：

75 4
einsum: 0.067826496087946
KDTree: 0.12196151306852698
500 4
einsum: 3.1056990439537913
KDTree: 0.85108971898444

我们可以看到，在较小的问题规模下，直接方法(einsum(更快，而对于较大的问题大小，KDTree获胜。