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

Question

假设我有一个m×n阶矩阵A和一个m×1阶向量。我想通过使用向量的元素作为每一行中的偏移量来从矩阵的每一行中提取元素 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 range(len(y)):
    A[i, y[i]] = #perform operations here

不使用任何显式循环。

预期输出是，

[3, 2, 5, 1, 1]

我正在使用 Python 和 NumPy 库。

Answer 1

以下遵循传统的 python 风格，但并未针对速度进行优化。试试看它是否适合您的用例。

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

b = [0, 1, 2, 3, 4]

result = [A[i][b_i] for i, b_i in enumerate(b)]
# -> [1, 2, 0, 1, 5]

Answer 2

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

In [270]: import numpy as np

In [271]: 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]])

In [272]: cols = [4, 2, 1, 3, 2]

之后，没有什么能阻止您使用 advanced indexing:

In [273]: rows = np.arange(A.shape[0])

In [274]: rows
Out[274]: array([0, 1, 2, 3, 4])

In [275]: A[rows, cols]
Out[275]: array([3, 2, 5, 1, 1])

In [276]: A[rows, cols] = -99

In [277]: A
Out[277]: 
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]])

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

How to extract elements from a matrix using a vector of indices?

python

numpy

vector

matrix

matrix-indexing