假设我有一个二维的numpy数组：
mat=[[5,5,3,6,3],
[3,2,7,8,1],
[7,5,5,2,0]]
和索引向量：
vec=[3,1,2]
我需要从相应的索引中提取3个元素。例如，向量中与矩阵第一行相对应的第一个元素是3。因此，我需要从第一行的索引3（从0开始）中取3个元素，即6。所以我需要的是[6,3,None]。
最终输出应为：
[[6,3,None],
[2,7,8],
[5,2,0]]
我尝试take使用花哨的索引，但都无法获得想要的输出。
任何帮助，将不胜感激！
你可以这样做 -
import numpy as np
mat=np.array([[5,5,3,6,3],
[3,2,7,8,1],
[7,5,5,2,0]])
mat = np.hstack((mat, np.ones((3,3))*np.nan))
vec=np.array([3,1,2])
idx = vec[:, None] + np.arange(0, 3)
print(mat[np.arange(3)[:,None], idx])
给-
[[ 6. 3. nan]
[ 2. 7. 8.]
[ 5. 2. 0.]]
首先，只需在原始数组上附加三列inf或None或某物。然后，vec通过添加从0开始的连续整数，从中创建一个2d索引数组，并简单地索引原始矩阵。
本文收集自互联网，转载请注明来源。
如有侵权，请联系[email protected] 删除。
我来说两句