python - python2中numpy的切片操作

Question

如果我有一个numpy数组X:

array([[ 0.13263767,  0.23149757,  0.57097612],
       [ 0.49629958,  0.67507182,  0.6758823 ]])

还有一个索引数组Y:

array([1, 2, 1])

我可以使用X[0:,Y]来索引X的第一行，它将输出:

array([ 0.23149757,  0.57097612,  0.23149757])

我的问题是，如果我有一个二维索引数组Z:

array([[1, 2, 1],
       [0, 1, 2]])

我想使用Z的第一行来索引X的第一行，使用Z的第二行来索引X 的第二行(Z 和 X 具有相同的行)。 所以一种方法是使用如下命令:

Row_0 = X[0:, Z[0]]
Row_1 = X[1:, Z[1]]

我想知道是否有一种简单的方法可以做到这一点。 谢谢

Answer 1

您可以使用花哨的索引来实现这一点:

>>> X[[[0], [1]], Z]
array([[ 0.23149757,  0.57097612,  0.23149757],
       [ 0.49629958,  0.67507182,  0.6758823 ]])

诀窍是索引第一个维度的数组必须与索引第二个维度的数组一起广播。在这种情况下:

>>> np.array([[0], [1]]).shape
(2, 1)
>>> Z.shape
(2, 3)

因此，返回的结果将是广播形状，(2, 3)，其索引取自第一个数组的第一维和第二个数组的第二维。

对于更一般的情况，您可以获得与以下相同的结果:

>>> X[np.arange(Z.shape[0])[:, None], Z]
array([[ 0.23149757,  0.57097612,  0.23149757],
       [ 0.49629958,  0.67507182,  0.6758823 ]])