在 numpy 中，如何使用索引列表從二維數組的第二維進行切片？

Question

最好用一個例子來解釋。 下面的代碼給了我想要的結果，但我想避免迭代/列表理解。

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]

我不想select 每行的相同列集，例如foo[:, [1, 3, 4]] 。

明確地說：是否有 Numpy function 用於此？ np.ix_似乎很接近，但似乎 select 完整列。 理想情況下，在示例代碼中沒有zip的情況下，我可以將[1, 3, 4]作為輸入。

我正在使用 Python 3.9 和 NumPy 1.19，盡管這無關緊要。 另外，如果有人可以建議一個更好的標題...哎呀

謝謝！

Answer 1

只需這樣做：

foo[zip(*indices)]

Answer 2

用這個：

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

Answer 3

這是一個純 Numpy oneliner：

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