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

【问题标题】：Iterating operation with two arrays using numpy使用 numpy 对两个数组进行迭代操作
【发布时间】：2018-11-17 22:53:16
【问题描述】：

我正在使用两个不同的阵列 (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))

【问题讨论】：

有三种基本选择：（1）使用嵌套循环，放弃numpy的所有速度优势； (2) 建立一个产品数组，这样你就可以做一个巨大的元素操作； (3) 将外循环保留为循环，但将内循环变为元素操作。您通常想做#3，因为它几乎具有#2 的所有速度优势，而没有内存成本。所以你只需要弄清楚如何从 A 中的一行对整个数组 B 进行操作，然后在 A 上循环这样做。
但是同时，两个 75x4 数组非常小，所以你可以写任何你觉得写得最舒服的东西，然后确保它不会太慢。

标签： python arrays numpy iteration

【解决方案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 胜出。

【讨论】：