Numpy: выберите значение в конкретной строке для каждого столбца матрицы

Question

У меня есть матрица 2D X = ((a11, a12, .. a1n), (a21 .. a2n) .. (am1, .. amn)) и вектор 1D y = [y1, ..., yn] каждый yi находится между 1 и m. Для каждого столбца i из X Я хочу выделить элемент в строке yi. То есть, я хочу выделить вектор z = (a_(y1 1), ... a_(yn n)).

Есть ли векторизованный способ сделать это?

Answer 1

Как об этом:

In [39]: x = np.arange(12).reshape(4,3) 

In [40]: y = np.array([0,3,2]) 

In [41]: x[y[None, :], np.arange(len(y))[None,:]][0] 
Out[41]: array([ 0, 10, 8]) 

In [42]: x 
Out[42]: 
array([[ 0, 1, 2], 
     [ 3, 4, 5], 
     [ 6, 7, 8], 
     [ 9, 10, 11]]) 

Answer 2

В качестве альтернативного решения, np.choose полезно для изготовления выбор.

>>> x = np.arange(16).reshape(4,4) 

Так x выглядит следующим образом:

array([[ 0, 1, 2, 3], 
     [ 4, 5, 6, 7], 
     [ 8, 9, 10, 11], 
     [12, 13, 14, 15]]) 

Теперь выбор значения в конкретной строке y в каждом столбце можно сделать так:

>>> y = np.array([3, 0, 2, 1]) 
>>> np.choose(y, x) 
array([12, 1, 10, 7])