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

Question

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

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

索引向量告诉我，我需要第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

由于您正在使用二维矩阵，因此建议您使用r(X)。然后，就您而言，您可以只使用np.diag：

a

但是，@ jpp的解决方案更具通用性。对于您的情况，我的解决方案很有用，因为您确实只想要矩阵的对角线。

Answer 3

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