使用numpy

Question

我正在使用两个不同的阵列（75x4），并且我在两个阵列之间应用了最短距离算法。

所以我想：

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

我怎么会用numpy来做这个？

基本上我想在数组2的每一行上的一行数组1之间执行一个操作，找到最小值，并将其存储在一个新数组中。然后为数组1的第2行做同样的事情，依此类推数组1的所有75行。

以下是公式I使用的代码。我得到的只是阵列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获胜。