Question

我有一个值矩阵和索引数组，我想遍历矩阵的每一行，只取与index数组中索引对应的值。

因此对于矩阵中的第1行，我取与数组中第一个元素对应的值。

Mat = [[0. 4. 4. 4.]
      [2. 8. 8. 2.]
      [7. 7. 7. 7.]
      ...
      [5. 5. 5. 5.]
      [7. 7. 7. 7.]
      [2. 6. 6. 6.]]

arr = [0 0 3 ... 0 0 3]

所以我想要第一行的第0个值，第二行的第0个，依此类推，我希望将它保存在最后的矢量中。

np.take(Mat,arr,axis=1,out=y)

但这不起作用，我得到一个矩阵作为回报而不是矢量。任何帮助将不胜感激。

由于

Answer 1

您可以对行和列使用integer array indexing：

res = mat[np.arange(len(mat)), arr]

这是一个演示：

mat = np.array([[0, 4, 3, 1],
                [2, 8, 8, 2],
                [1, 2, 3, 4],
                [5, 4, 3, 2],
                [4, 3, 2, 1],
                [2, 5, 6, 4]])

arr = np.array([0, 0, 3, 0, 0, 3])

res = mat[np.arange(len(mat)), arr]

print(res)

array([0, 2, 4, 5, 4, 4])

Answer 2

这可以使用choose方法完成：

np.choose(indices, matrix.T)

从numpy矩阵中提取某些指数

2 个答案: