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)

输出：

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]

我不想想从每一行中选择相同的一组列，例如foo[:, [1, 3, 4]]。

明确地说：是否有 Numpy 函数？ np.ix_ 似乎很接近，但它似乎选择了完整的列。理想情况下，我可以将 [1, 3, 4] 作为输入提供给示例代码中没有 zip 的内容。

我使用的是 Python 3.9 和 NumPy 1.19，尽管这无关紧要。另外，如果有人可以提出更好的标题...... yikes

谢谢！

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]

在 numpy 中，如何使用索引列表从二维数组的第二维进行切片？

3 个答案: