按行索引NumPy数组

Question

说我有一个NumPy数组:

>>> X = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
>>> X
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12]])

以及我想为每一行选择的索引数组:

>>> ixs = np.array([[1, 3], [0, 1], [1, 2]])
>>> ixs
array([[1, 3],
       [0, 1],
       [1, 2]])

如何索引数组X,以便为XI中的每一行选择指定的两个索引ixs？

因此,对于这种情况,我想为第一行选择元素1和3,为第二行选择元素0和1,依此类推.输出应该是:

array([[2, 4],
       [5, 6],
       [10, 11]])

缓慢的解决方案是这样的:

output = np.array([row[ix] for row, ix in zip(X, ixs)])

但是对于极长的阵列来说,这可能会有点慢.如果没有使用NumPy的循环,有没有更快的方法来做到这一点？

编辑: 2.5K*1M阵列的一些非常近似的速度测试,2K宽ixs(10GB):

np.array([row[ix] for row, ix in zip(X, ixs)]) 0.16s

X[np.arange(len(ixs)), ixs.T].T 0.175s

X.take(idx+np.arange(0, X.shape[0]*X.shape[1], X.shape[1])[:,None]) 33S

np.fromiter((X[i, j] for i, row in enumerate(ixs) for j in row), dtype=X.dtype).reshape(ixs.shape) 2.4S

Answer 1

你可以用这个:

X[np.arange(len(ixs)), ixs.T].T

这是复杂索引的参考.