python，如何从矩阵的每一列中选择元素

Question

我需要根据索引向量从矩阵的每一列中提取一个元素。说：

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

Index vector告诉我需要第1列的第一个元素，第2列的第二个元素，第3列的第三个元素。你知道吗

输出应该是[1,5,8]。没有显式循环，我怎么写出来？你知道吗

谢谢

Answer 1

您可以使用advanced indexing：

index = np.array([0,1,2])
matrix = np.array([[1,4,7],[2,5,8],[3,6,9]])

res = matrix[np.arange(matrix.shape[0]), index]
# array([1, 5, 9])

对于第二个示例，请反转索引：

index = np.array([0,1,1])
matrix = np.array([[1,4,7],[2,5,8],[3,6,9]])

res = matrix[index, np.arange(matrix.shape[1])]
# array([1, 5, 8])

Answer 2

val = [matrix[i][index[i]] for i in range(0, len(index))]

Answer 3

既然您使用的是二维矩阵，我建议您使用numpy。那么，在您的例子中，您可以使用^{}：

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

然而，@jpp的解决方案更具普遍性。我的解在你的例子中很有用，因为你只需要矩阵的对角线。你知道吗