使用 3x3 索引数组广播 3x3 数组

Question

我有一个名为 data 的 3x3 数组和一个名为 idx 的 3x3 索引数组。我希望能够使用广播在 idx 给出的索引处获得一个由 data 组成的新数组。我可以天真地处理这个问题并在 for 循环中执行，如下例所示，然后将其与暴力破解的 expected 数组进行比较：

import numpy as np
data = np.array([[0.5, 1.5, 2.5], [0.5, 1.5, 2.5], [0.5, 1.5, 2.5]])
idx = np.array([[0,-1,-2], [1,0,-1], [2,1,0]])
expected = np.array([[0.5, 2.5, 1.5], [1.5, 0.5, 2.5], [2.5, 1.5, 0.5]])

result = np.zeros(np.shape(data))

for i in range(len(idx)):
    for j in range(len(idx[i])):
        result[i,j]=data[i, idx[i,j]]
        
print(expected==result)
# Gives: 3x3 array of True

我之所以把它带到这里，是因为我需要将它应用到一个 NxM 数组，如果我像上面的例子那样应用它，这将花费很长时间来计算。

我发现了两个与我的问题相关的类似问题 ( and )，但我不确定如何将其应用于任意大的二维数组。我尝试了以下但没有成功：

result = data[np.ix_(*idx)] 
# Gives Error: too many indices for array: array is 2-dimensional, but 3 were indexed

和

for i in range(len(idx)):
    sub = np.ix_(idx[i])
    print(sub)
    # Gives: (array([ 0, -1, -2]),)
    result[i] = data[sub] 
print(result)
# Gives Error: could not broadcast input array from shape (3,3) into shape (3,)

必须有一种方法可以简单地使用 Numpy 来做到这一点，但我还没有找到。

Answer 1

如果还明确指定列值，您将获得该行为

import numpy as np
data = np.array([[0.5, 1.5, 2.5], [0.5, 1.5, 2.5]])
idx = np.array([[0,-1,-2], [1,0,-1]])
expected = np.array([[0.5, 2.5, 1.5], [1.5, 0.5, 2.5]])
print(data[np.arange(len(data)).reshape(-1,1),idx] == expected)

输出：

[[ True  True  True]
 [ True  True  True]]

使用 3x3 索引数组广播 3x3 数组

Broadcasting 3x3 array with 3x3 array of indices

python

arrays

numpy

array-broadcasting