python一维数组作为列下标分别获取二维数组每一行的对应的元素简单实现

1. 代码实现

import numpy as np

i=np.array([1,3,2])     #i=[1,3,2]也可以
y_hat=np.array([[0.3,0.6,0,0.1,0],
　　　　　　　　　　[0.8,0,0,0,0.2],
　　　　　　　　　　[0.1,0.1,0.2,0.6,0]])
result=y_hat[np.arange(0,len(y_hat),1),i]    #np.arange(0,len(y_hat),1)为[0,...]列表
print(result)     #结果为[0.6 0.  0.2]