你可以使用
NumPy's advanced indexing
A[np.arange(A.shape[0])[:,None],B]
一个人也可以使用
linear indexing
m,n = A.shapeout = np.take(A,B + n*np.arange(m)[:,None])
样品运行
In [40]: AOut[40]: array([[2, 4, 5, 3], [1, 6, 8, 9], [8, 7, 0, 2]])In [41]: BOut[41]: array([[0, 0, 1, 2], [0, 3, 2, 1], [3, 2, 1, 0]])In [42]: A[np.arange(A.shape[0])[:,None],B]Out[42]: array([[2, 2, 4, 5], [1, 9, 8, 6], [2, 0, 7, 8]])In [43]: m,n = A.shapeIn [44]: np.take(A,B + n*np.arange(m)[:,None])Out[44]: array([[2, 2, 4, 5], [1, 9, 8, 6], [2, 0, 7, 8]])
