python, how to select element from each column of matrix

Question

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

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

Answer 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.

Answer 3

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