Python, как выбрать элемент из каждого столбца матрицы

Question

Мне нужно извлечь один элемент из каждого столбца матрицы в соответствии с индексом вектора. Скажи:

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

Индексный вектор говорит мне, что мне нужен первый элемент из столбца 1, второй элемент из столбца 2 и третий элемент из столбца 3.

Вывод должен быть [1,5,8]. Как я могу выписать это без явного цикла?

Спасибо

Answer 1

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

Answer 2

Вы можете использовать расширенное индексирование :

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 3

Поскольку вы работаете с 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 более обобщенно. Мое решение полезно в вашем случае, потому что вы действительно просто хотите диагональ матрицы.