在 numpy 中创建索引数组 - 消除双 for 循环答案

【问题标题】：Creating index array in numpy - eliminating double for loop在 numpy 中创建索引数组 - 消除双 for 循环
【发布时间】：2015-03-15 02:21:22
【问题描述】：

我有一些物理模拟代码，用 python 编写并使用 numpy/scipy。分析代码显示 38% 的 CPU 时间花在一个双重嵌套的 for 循环中 - 这似乎过多，所以我一直在尝试减少它。

循环的目标是创建一个索引数组，显示一维数组的哪些元素与二维数组的元素相等。

indices[i,j] = where(1D_array == 2D_array[i,j])

例如，如果1D_array = [7.2, 2.5, 3.9] 和

2D_array = [[7.2, 2.5] 
            [3.9, 7.2]]

我们应该有

indices = [[0, 1]
           [2, 0]]

我目前将其实现为

for i in range(ni):
    for j in range(nj):
        out[i, j] = (1D_array - 2D_array[i, j]).argmin()

argmin 是需要的，因为我正在处理浮点数，因此等式不一定准确。我知道一维数组中的每个数字都是唯一的，并且二维数组中的每个元素都有一个匹配项，所以这种方法给出了正确的结果。

有没有办法消除双重for循环？

注意：

我需要索引数组来执行以下操作：

f = complex_function(1D_array)
output = f[indices]

这比替代方法更快，因为 2D 数组的大小为 NxN，而 1D 数组的大小为 1xN，并且 2D 数组有许多重复值。如果有人可以提出一种不同的方式来获得相同的输出而不通过索引数组，那也可能是一个解决方案

【问题讨论】：

1D_array 总是排序的？
@AshwiniChaudhary，不，不是。事实上，它永远不会。我将编辑示例以删除它。
为此，我认为 1D_array 中的条目不会重复。为什么不用 1D_array 制作字典，将值作为键，将索引作为值？ IE。 {0:7.2, 1:2.5, 2:3.9} 然后你只需要将 dict 应用于数组。

标签： python arrays performance numpy indexing

【解决方案1】：

在纯 Python 中，您可以在 O(N) 时间使用字典来执行此操作，唯一的时间损失将是涉及的 Python 循环：

>>> arr1 = np.array([7.2, 2.5, 3.9])
>>> arr2 = np.array([[7.2, 2.5], [3.9, 7.2]])
>>> indices = dict(np.hstack((arr1[:, None], np.arange(3)[:, None])))
>>> np.fromiter((indices[item] for item in arr2.ravel()), dtype=arr2.dtype).reshape(arr2.shape)
array([[ 0.,  1.],
       [ 2.,  0.]])

【讨论】：

【解决方案2】：

其他人建议的字典方法可能有效，但它要求您提前知道目标数组（二维数组）中的每个元素在搜索数组（一维数组）中都有完全匹配。即使原则上应该是这样，您仍然必须处理浮点精度问题，例如试试这个.1 * 3 == .3。

另一种方法是使用 numpy 的 searchsorted 函数。 searchsorted 采用已排序的一维搜索数组和任何 traget 数组，然后为目标数组中的每个项目在搜索数组中找到最接近的元素。我已根据您的情况调整了此 answer，请查看它以了解 find_closest 函数的工作原理。

import numpy as np

def find_closest(A, target):
    order = A.argsort()
    A = A[order]

    idx = A.searchsorted(target)
    idx = np.clip(idx, 1, len(A)-1)
    left = A[idx-1]
    right = A[idx]
    idx -= target - left < right - target
    return order[idx]

array1d = np.array([7.2, 2.5, 3.9])
array2d = np.array([[7.2, 2.5],
                    [3.9, 7.2]])

indices = find_closest(array1d, array2d)
print(indices)
# [[0 1]
#  [2 0]]

【讨论】：

【解决方案3】：

要摆脱两个 Python for 循环，您可以通过向数组添加新轴（使它们彼此可广播）“一次性”完成所有相等比较。

请记住，这会生成一个包含 len(arr1)*len(arr2) 值的新数组。如果这是一个非常大的数字，那么这种方法可能不可行，具体取决于您的内存限制。否则，它应该相当快：

>>> (arr1[:,np.newaxis] == arr2[:,np.newaxis]).argmax(axis=1)
array([[0, 1],
       [2, 0]], dtype=int32)

如果您需要获取arr1 中最接近 匹配值的索引，请使用：

np.abs(arr1[:,np.newaxis] - arr2[:,np.newaxis]).argmin(axis=1)

【讨论】：