如果我有一个 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])
我的问题是,如果我有一个2维的索引数组 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]]
我想知道是否有一种简单的方法可以做到这一点。 谢谢
答案 0 :(得分:3)
您可以使用花式索引来实现:
>>> 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 ]])