python, how to select element from each column of matrix

问题:

I need to extract one element from each column of a matrix according to an index vector. Say:

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

Index vector tells me I need the first element from column 1, second element from column 2, and third element from column 3.

The output should be [1,5,8]. How can I write it out without explicit loop?

Thanks

回答1:

You can use 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]) 

For your second example, reverse your indices:

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]) 

回答2:

Since you're working with 2-dimensional matrices, I'd suggest using numpy. Then, in your case, you can just use 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]) 

However, @jpp's solution is more generalizable. My solution is useful in your case because you really just want the diagonal of your matrix.

回答3:

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

文章来源: python, how to select element from each column of matrix