python - 使用 numpy 中的数组清理数组索引

Question

我一次又一次地遇到这个问题，并且似乎无法找到一个干净的解决方案。所以我试图用另一个数组来索引一个数组。我有一个二维 numpy 数组。一个 1d numpy 数组的长度与我尝试索引的 2d 数组的第一个维度相同，元素代表我尝试提取的列的索引:

import numpy as np

A = np.random.rand((5,3))

B = np.asarray([2,1,2,0,1])

我想要的行为是为所有行提取数组 B 中的相应列。这可以通过

来完成

C = A[np.arange(A.shape[0]),B]

但我可以想象有更好的方法来实现这种行为。使用 : 作为第一行的索引会产生错误的行为。

如果有一种更简洁的方法来做到这一点那就太好了。我真的习惯了 Matlab 中的这个数组索引，但也许 numpy 中没有等效的东西。使用 bool 索引当然是一种选择，但这也需要始终转换数组。

Answer 1

我认为你关心的是np.choose(B,A.T):

In [125]: A
Out[125]: 
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11],
       [12, 13, 14]])

In [126]: B = np.asarray([2,1,2,0,1])

In [127]: np.choose(B,A.T)
Out[127]: array([ 2,  4,  8,  9, 13])