在Numpy中逐行索引

时间:2013-11-04 19:48:41

标签: python numpy indexing

我有两个矩阵,AB

A = array([[2., 13., 25., 1.], [ 18., 5., 1., 25.]])
B = array([[2, 1], [0, 3]])

我想用A的每一行为B的每一行编制索引,从而产生切片:

array([[25., 13.], [18., 25.]])

也就是说,我基本上想要这样的东西:

array([A[i,b] for i,b in enumerate(B)])

有没有办法直接指出这个?我能做的最好的就是这个“平板黑客”:

A.flat[B + arange(0,A.size,A.shape[1])[:,None]]

2 个答案:

答案 0 :(得分:9)

@ Ophion的答案很棒,值得赞扬,但我想补充一些解释,并提供更直观的结构。

不是旋转B然后再旋转结果,最好只旋转arange。我认为这提供了最直观的解决方案,即使它需要更多的字符:

A[((0,),(1,)), B]

或等效

A[np.arange(2)[:, None], B]

这是有效的,因为这里真正发生的是,您正在创建一个i数组和一个j数组,每个数组的形状都与您想要的结果相同。

i = np.array([[0, 0],
              [1, 1]])
j = B

但你可以只使用

i = np.array([[0],
              [1]])

因为它会广播以匹配B(这是np.arange(2)[:,None]给出的)。

最后,为了使其更加通用(不知道2arange大小),您还可以使用

i生成B
i = np.indices(B.shape)[0]

然而,您构建ij,您只需将其称为

>>> A[i, j]
array([[ 25.,  13.],
       [ 18.,  25.]])

答案 1 :(得分:7)

不漂亮,但是:

A[np.arange(2),B.T].T
array([[ 25.,  13.],
       [ 18.,  25.]])