Question

我有一些物理模拟代码，用python编写并使用numpy / scipy。对代码进行分析表明，38％的CPU时间花费在一个双重嵌套的for循环中 - 这似乎过多，所以我一直试图将其削减。

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

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，因为我处理浮点数，所以相等不一定精确。我知道1D数组中的每个数字都是唯一的，并且2D数组中的每个元素都有匹配，因此这种方法可以得到正确的结果。

有没有办法消除双循环？

注意：

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

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

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

Answer 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.]])

Answer 2

要摆脱两个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)

Answer 3

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

另一种方法是使用numpy的searchsorted函数。 searchsorted采用排序的1d搜索数组，然后任何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]]

在numpy中创建索引数组 - 消除double for循环

3 个答案: