How to index a numpy array of dimension N with a 1-dimensional array of shape (N,)

Question

I would like to index an array of dimension N using an array of size (N,). For example, let us consider a case where N is 2.

import numpy as np
foo = np.arange(9).reshape(3,3)
bar = np.array((2,1))

>>> foo
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])

>>>bar
array([2, 1])

>>>foo[bar[0],bar[1]]
7

This works fine. However, with this method, I would need to write N times bar[i], which is not a nice solution if N is high.

The following command does not give the result that I need:

>>>foo[bar]
array([[6, 7, 8],
       [3, 4, 5]])

What could I do to get the result that I want in a nice and concise way?

Answer 1

I think you can turn bar into tuple:

foo[tuple(bar)]
# 7