Numpy 二维数组 - 从指定索引中获取 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] 。

最终的 output 应该是：

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

我尝试使用take和花哨的索引，但无法获得所需的 output。

任何帮助，将不胜感激！

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.]]

首先只是 append 具有三列inf或None或其他东西的原始数组。 然后通过添加从 0 开始的连续整数从vec创建一个二维索引数组，并简单地索引原始矩阵。