如何使用numpy使用索引向量从矩阵中提取元素？

Question

假设我有一个mxn阶的矩阵A和mx1的矢量.我想通过使用向量的元素作为每行中的偏移量来从矩阵A的每一行中提取元素.

例如,

A = [[3, 0, 0, 8, 3],
    [9, 3, 2, 2, 6],
    [5, 5, 4, 2, 8],
    [3, 8, 7, 1, 2],
    [3, 9, 1, 5, 5]]

和一个矢量

y = [4,2,1,3,2]

我想要实现的是我想要一种方法来提取A的元素,使得向量的每个元素索引A的相应行中的元素,即实现

for i in xrange(y):
    A[i,y[i]] = #perform operations here

不使用任何显式循环.

预期的输出是,

[3,2,5,1,1]

我正在使用Python和numpy库.

谢谢您的帮助!

Answer 1

您应该首先将列表A转换为NumPy数组:

>>> import numpy as np
>>> A = np.array([[3, 0, 0, 8, 3],
...               [9, 3, 2, 2, 6],
...               [5, 5, 4, 2, 8],
...               [3, 8, 7, 1, 2],
...               [3, 9, 1, 5, 5]])
...
>>> y = [4, 2, 1, 3, 2]

之后,没有什么能阻止您使用高级索引:

>>> A[np.arange(A.shape[0]), y]
array([3, 2, 5, 1, 1])
>>> A[np.arange(A.shape[0]), y] = -99
>>> A
array([[  3,   0,   0,   8, -99],
       [  9,   3, -99,   2,   6],
       [  5, -99,   4,   2,   8],
       [  3,   8,   7, -99,   2],
       [  3,   9, -99,   5,   5]])