我有一些物理仿真代码,使用python和numpy / scipy编写.对代码进行性能分析表明,38％的CPU时间都花在了一个双重嵌套的for循环中-这似乎过多,因此我一直在尝试减少它.

循环的目标是创建一个索引数组,以显示2D数组的元素等于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,因此等式不一定精确.我知道一维数组中的每个数字都是唯一的,而二维数组中的每个元素都有一个匹配项,因此这种方法可以得出正确的结果.

有什么办法可以消除double for循环吗？

注意：

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

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

这比替代方法更快,因为2D数组的大小为NxN,而1D数组的大小为1xN,并且2D数组具有许多重复值.如果有人可以建议不通过索引数组而达到相同输出的另一种方式,那也可能是一个解决方案

解决方法:

其他人建议的字典方法可能可行,但它要求您提前知道目标数组(2d数组)中的每个元素在搜索数组(1d数组)中都有完全匹配的内容.即使原则上是正确的,您仍然必须处理浮点精度问题,例如,尝试使用.1 * 3 == 0.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]]