Numpy матрица: выберите столбец на основе списка

Question

У меня есть массив numpy и список, в котором перечислены строки, которые я хочу выбрать. Каков наилучший способ выполнить эту операцию?

import numpy as np

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

b = np.array([[1],
              [0],
              [2]])

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

np.array([[2],
         [4],
         [9]])

Я пробовал np.take (), но это не работает.

С уважением

РЕДАКТИРОВАТЬ: так как это необходимо делать повторно для большого массива, я ищу векторизованный подход (без циклов)

Answer 1

Если вы удалите посторонние размеры из b

b = np.sqeeze(b)

Вы можете использовать следующее:

a[np.arange(len(b)), b]

Answer 2

Последние версии добавили take_along_axis, который делает то, что вы хотите:

In [96]: a = np.array([[1,2,3], 
    ...:               [4,5,6], 
    ...:               [7,8,9]]) 
    ...:  
    ...: b = np.array([[1], 
    ...:               [0], 
    ...:               [2]])                                                                           
In [97]: np.take_along_axis(a,b,axis=1)                                                                
Out[97]: 
array([[2],
       [4],
       [9]])

Он работает так же, как @ Nils answer , a[np.arange(3), np.squeeze(b)], но лучше обрабатывает размеры.

Подобные недавние вопросы:

Эффективное индексирование матрицы на столбец в numpy

сохраняет элементы np.ndarray значениями другой массив np.array (векторизация)

Answer 3

Это не очень Pythoni c, но это должно помочь вашей проблеме:

res = np.zeros(len(b))
for i, row in enumerate(a):
    res[i] = row[b[i]]

print(res)

то же самое в одной строке:

a[[i[0] for i in b],[i for i in range(len(b))]]

Answer 4

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

np.array([a[i,b[i]] for i in range(len(b))]

Answer 5

Примерно так:

arr = []
In [892]: for c,i in enumerate(a): 
     ...:      arr.append(a[c][b[c]]) 
     ...:                

In [895]: arr = np.array(arr)                                                                                                                                                                               

In [896]: arr                                                                                                                                                                                               
Out[896]: 
array([[2],
       [4],
       [9]])