Indexing a multi-dimensional numpy array with an index array [duplicate]

Question

I am having a hard time indexing a two-dimensional array given an indexing array.

Suppose I have:

# Indexing array 
idx = np.array([0,2,1,2,1])

# Array to be indexed
my_array = np.array([[0,2,1], [0,5,1], [1,2,1], [5,1,3], [2,6,2]])

And the expected output should be the 0 index of the first entry of my_array, the 2 index of the second entry of my_array and so on, thus:

# Expected output
expected_array = np.array([0,1,2,3,6])

I have actually made it using a for-loop:

# Using for-loop
expected_array = np.array([])
for i in range(len(my_array)):
    expected_array = np.append(idx, my_array[i][idx[i]])

But I am wondering if there is a way of possibly doing array-indexing without the use of for-loops? Assume that len(idx)==len(my_array) all the time.

Cory Kramer · Accepted Answer · 2017-07-24 15:00:58Z

3

You can use numpy.arange to make an array to index rows, which will just be each row at a time. Then use your idx array to index out of each row.

>>> my_array[np.arange(my_array.shape[0]), idx]
array([0, 1, 2, 3, 6])

answered Jul 24, 2017 at 15:00

Cory Kramer

119k19 gold badges176 silver badges233 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Indexing a multi-dimensional numpy array with an index array [duplicate]

1 Answer 1

Comments

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Linked

Related