Broadcasting 3x3 array with 3x3 array of indices-CodePudding

I have a 3x3 array called data, and a 3x3 array of indices called idx. I'd like to be able to use broadcasting to get a new array composed of data at the indices given by idx. I can approach this naively and do it in a for-loop as in the example below, then compare it to the brute-forced expected array:

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

The reason why I am bringing this here, is because I need to apply this to an NxM array, which will take a long time to compute if I apply it as in the example above.

I found two similar questions (one and two) which are related to my question, but I am not sure how to apply it to an arbitrarily large 2D array. I have attempted the following with no luck:

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

and

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,)

There has to be a way to do this simply with Numpy that I just haven't found.

CodePudding user response：

You will get that behavior if also explicitly specify the column values

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)

Output:

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