Как извлечь элементы из матрицы с помощью вектора индексов?

Предположим, у меня есть матрица A порядка m × n и вектор порядка m × 1. Я хотел бы извлечь элементы из каждой строки матрицы A, используя элементы вектора в качестве смещения в каждой строке.

Например,

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

и вектор

y = [4, 2, 1, 3, 2]

Я хочу добиться такого способа извлечения элементов A, чтобы каждый элемент вектора индексировал элемент в соответствующей строке A, т. Е. Реализуя

for i in range(len(y)):
    A[i, y[i]] = #perform operations here

без использования каких-либо явных циклов.

Ожидаемый результат:

[3, 2, 5, 1, 1]

Я использую Python и библиотеку NumPy.


person Suyash Shetty    schedule 18.05.2016    source источник
comment
Не могли бы вы также добавить ожидаемый результат?   -  person AKS    schedule 18.05.2016


Ответы (2)


Вам следует начать с преобразования списка A в массив NumPy:

>>> import numpy as np
>>> A = np.array([[3, 0, 0, 8, 3],
...               [9, 3, 2, 2, 6],
...               [5, 5, 4, 2, 8],
...               [3, 8, 7, 1, 2],
...               [3, 9, 1, 5, 5]])
...
>>> y = [4, 2, 1, 3, 2]

И после этого вам ничто не мешает использовать расширенное индексирование :

>>> A[np.arange(A.shape[0]), y]
array([3, 2, 5, 1, 1])
>>> A[np.arange(A.shape[0]), y] = -99
>>> A
array([[  3,   0,   0,   8, -99],
       [  9,   3, -99,   2,   6],
       [  5, -99,   4,   2,   8],
       [  3,   8,   7, -99,   2],
       [  3,   9, -99,   5,   5]])
person Tonechas    schedule 18.05.2016

Следующее соответствует обычному стилю Python, но не оптимизировано для скорости. Попробуйте и посмотрите, подходит ли он для вашего варианта использования.

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

b = [0, 1, 2, 3, 4]

result = [A[i][b_i] for i, b_i in enumerate(b)]
# -> [1, 2, 0, 1, 5]
person Razzi Abuissa    schedule 18.05.2016