python - 在 numpy 中创建索引数组 - 消除双循环

Question

我有一些物理模拟代码，用 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 是必需的，因为我正在处理 float ，因此相等性不一定准确。我知道一维数组中的每个数字都是唯一的，而二维数组中的每个元素都有一个匹配项，所以这种方法给出了正确的结果。

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

注意:

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

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

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

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