python，如何從矩陣的每一列中選擇元素

Question

我需要根據索引向量從矩陣的每一列中提取一個元素。 說：

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

索引向量告訴我，我需要第1列的第一個元素，第二列的第2個元素，第3列的第三個元素。

輸出應為[1,5,8] 。 如何在沒有顯式循環的情況下將其寫出？

謝謝

Answer 1

您可以使用高級索引 ：

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

由於您使用的是二維矩陣，因此建議您使用numpy 。 然后，根據您的情況，您可以使用np.diag ：

>>> 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的解決方案更具通用性。 我的解決方案在您的情況下很有用，因為您確實只想要矩陣的對角線。

Answer 3

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