Матричное индексирование, извлечение разных столбцов для каждой строки - PullRequest
0 голосов
/ 16 сентября 2018

Пример с маленькими массивами внизу. У меня есть массив данных размером MxN (N> 100) и меньший массив A с размерами Mx100.

Массив A содержит столбец с указанием массива данных.

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

Я пытался

 data[:,A]

но это возвращает мне массив MxMx100. Я просто хочу массив Mx100. Я мог бы сделать это в цикле for, но это занимает очень много времени. Как я могу сделать это правильно?

Пример:

data=np.array([[1,2,3],[4,5,6],[7,8,9]])
array([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
A=np.array([[0,1],[1,2],[0,2]])

Я хочу извлечь: первая строка: индексы 0 и 1. вторая строка: индексы 1 и 2. третья строка: индексы 0 и 2

Желаемый результат:

array([[1,2],
[5,6],
[7,9]])

1 Ответ

0 голосов
/ 17 сентября 2018

Использование broadcasting

>>> data[np.arange(A.shape[0])[:,None], A]

array([[1, 2],
       [5, 6],
       [7, 9]])
...