In numpy, how do I use a list of indices to slice from the second dimension of a 2D array?

Question

Best explained with an example. The following code gives me the result that I want, but I want to avoid the iteration/list comprehension.

import numpy as np

foo = np.array([range(100, 105), range(200, 205), range(300, 305)])
print("Data:\n", foo)

# "Column 1 of row 0, column 3 of row 1, ..."
indices = list(zip(range(len(foo)), np.array([1, 3, 4])))
print("Indices to select from foo:\n", indices)

# This works, but surely there's a better way?
values = np.array([foo[row, col] for row, col in indices])
print("Value for the given column of each row:\n", values)

Output:

Data:
 [[100 101 102 103 104]
 [200 201 202 203 204]
 [300 301 302 303 304]]
Indices to select from foo:
 [(0, 1), (1, 3), (2, 4)]
Value for the given column of each row:
 [101 203 304]

I do not want to select the same set of columns from each row, e.g. foo[:, [1, 3, 4]].

To put it explicitly: is there a Numpy function for this? np.ix_ seems close, but it seems to select full columns. Ideally I would be able to give [1, 3, 4] as an input to something without the zip in the example code.

I'm using Python 3.9 and NumPy 1.19, though it shouldn't matter. Also, if anyone can suggest a better title... yikes

Thanks!

Answer 1

simply do this:

foo[zip(*indices)]

Answer 2

Use this:

>>> index = np.array([1, 3, 4])
>>> foo[range(foo.shape[0]), index]
array([101, 203, 304])

Answer 3

Here's a pure Numpy oneliner:

foo[np.arange(foo.shape[0]), indices]