CODEWITHSUNDEEP
pythonarraysnumpy
NightElfik
NightElfik

Reputation: 4537

How to index an np.array with a list of indices in Python

Suppose I have an N-dimensional np.array (or just a list) and a list of N indices. What is the preferred/efficient way to index the array without using loops?

# 4D array with shape of (2, 3, 4, 5)
arr = np.random.random((2, 3, 4, 5))
index = [0, 2, 1, 3]
result = ??? # Equivalent to arr[0, 2, 1, 3]

Additionally, supplying only a 3D index the result should be an array of the last dimension.

index = [0, 2, 1]
result2 = ??? # Equivalent to arr[0, 2, 1]

Please note that I am not able to just index with the usual syntax because the implementation has to handle arrays of different shapes.

I am aware that NumPy supports indexing by an array but that behaves differently as it cherry-picks values from the array rather by indexing by dimension (https://docs.scipy.org/doc/numpy/user/basics.indexing.html).

Upvotes: 1

Views: 297

Answers (1)

unutbu
unutbu

Reputation: 879471

Per the docs:

If one supplies to the index a tuple, the tuple will be interpreted as a list of indices.

Therefore, change index to a tuple:

In [46]: np.allclose(arr[tuple([0,2,1])], arr[0,2,1])
Out[46]: True

In [47]: np.allclose(arr[tuple([0,2,1,3])], arr[0,2,1,3])
Out[47]: True

Upvotes: 3

Related Questions