numpy ：取特定行的特定列

假如A是一个3*4的矩阵，要取第0行的第二列，第1行的第三列，第2行的第0列

A = np.random.randn(3,4)
B = [2,3,0]
A
Out[8]: 
array([[-0.75609069, -1.11338785,  0.67113714,  0.56159782],
       [ 0.0649341 , -1.94732805, -0.43665225,  1.85291592],
       [-0.66896335, -0.65058874, -0.51694697,  1.55056179]])

若直接使用A[:, B]是分别得到每行的0、2、3列

A[:,B]
Out[12]: 
array([[ 0.67113714,  0.56159782, -0.75609069],
       [-0.43665225,  1.85291592,  0.0649341 ],
       [-0.51694697,  1.55056179, -0.66896335]])

若要得到预期结果可使用A[np.arange(A.shape[0], B)]得到

A[np.arange(3),B]
Out[19]: array([ 0.67113714,  1.85291592, -0.66896335])