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

Question

假设我有一个 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,None]。

最终的输出应该是：

[[6,3,None],
 [2,7,8],
 [5,2,0]]

我尝试使用take 和精美的索引，但无法获得所需的输出。

任何帮助将不胜感激！

Answer 1

你可以这样做 -

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 开始添加连续整数并简单地索引原始矩阵。