Python Numpy 2d数组-从指定索引中获取N个元素

假设我有一个2d 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，无]

最终输出应为：

[[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添加顺序整数，并简单地为原始矩阵编制索引