CODEWITHSUNDEEP
pythonnumpy
Average Algorithm Guy
Average Algorithm Guy

Reputation: 23

Numpy array of multiple indices replace with a different matrix

I have an array of 2d indices.

indices = [[2,4], [6,77], [102,554]]

Now, I have a different 4-dimensional array, arr, and I want to only extract an array (it is an array, since it is 4-dimensional) with corresponding index in the indices array. It is equivalent to the following code.

for i in range(len(indices)):
    output[i] = arr[indices[i][0], indices[i][1]]

However, I realized that using explicit for-loop yields a slow result. Is there any built-in numpy API that I can utilized? At this point, I tried using np.choose, np.put, np.take, but did not succeed to yield what I wanted. Thank you!

Upvotes: 2

Views: 866

Answers (2)

Naveen Kumar reddy kambham
Naveen Kumar reddy kambham

Reputation: 274

Try using the take function of numpy arrays. Your code should be something like: 

outputarray= np.take(arr,indices)

Upvotes: 0

Divakar
Divakar

Reputation: 221684

We need to index into the first two axes with the two columns from indices (thinking of it as an array).

Thus, simply convert to array and index, like so -

indices_arr = np.array(indices)
out = arr[indices_arr[:,0], indices_arr[:,1]]

Or we could extract those directly without converting to array and then index -

d0,d1 = [i[0] for i in indices], [i[1] for i in indices]
out = arr[d0,d1]

Another way to extract the elements would be with conversion to tuple, like so -

out = arr[tuple(indices_arr.T)]

If indices is already an array, skip the conversion process and use indices in places where we had indices_arr.

Upvotes: 1

Related Questions