Let's say I have the numpy array x where x is

x = np.array(
  [
    [1, 2],
    [3, 4]
  ]
)

I also have a list of indices i where i = [0, 1]. I would like to get an array of the the value at index n[i] for every n in x.

The optimal output is

np.array([1, 4])

This can very obviously be done with a for loop, but I was wondering if there was a simpler *numpy* way to do it.

CodePudding user response：

You can use indexing:

x[np.arange(x.shape[0]), i]

output: array([1, 4])