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創建一個二維索引數組，並簡單地索引原始矩陣。