Извлечь массив значений с использованием индексов

Question

Мне нужно извлечь определенные значения из многомерного массива, которые не являются последующими.

import numpy as np  
A = np.array([[[ 0.,   4.,   0.  ], 
       [ 0.19230769, 4.03846154, 0.  ], 
       [-0.4,   4.8,   0.  ], 
       [ 2.,   1.,   0.  ]], 

       [[ 1.2,   3.4,   0.  ], 
       [ 2.11538462, 4.42307692, 0.  ], 
       [ 0.,   4.,   0.  ], 
       [ 3.6,   1.8,   0.  ]], 

       [[ 1.8,   3.1,   0.  ], 
       [ 3.17307692, 4.63461538, 0.  ], 
       [ 0.,   4.,   0.  ], 
       [ 4.,   2.,   0.  ]]]) 

Для каждого 4x3 блока Я хочу, чтобы извлечь произвольную строку

Например, следующие элементы:

A[0,2,:] 
A[1,1,:] 
A[2,1,:] 

Так-существу Ряды B = [2,1,1], который дал бы мне:

[-0.4   4.8   0.  ] 
[ 2.11538462 4.42307692 0.  ] 
[ 3.17307692 4.63461538 0.  ] 

Как это сделать эффективно?

Answer 1

Вы можете использовать "advanced indexing":

In [99]: A[[0,1,2], [2,1,1], :] 
Out[99]: 
array([[-0.4  , 4.8  , 0.  ], 
     [ 2.11538462, 4.42307692, 0.  ], 
     [ 3.17307692, 4.63461538, 0.  ]]) 

Здесь индексирование массивы

ind1 = [0, 1, 2] 
ind2 = [2, 1, 1] 

и так ind1 индексирует первую ось A и ind2 индексирует второй оси, а третья ось (с :), результирующий массив, result, имеет ту же форму, что и ind1 и ind2 - то есть (3,) - плюс форма полностью нарезанной оси, которая также равна (3,). Таким образом, result.shape является (3, 3) и

result[i, j] = A[ind1[i], ind2[i], j] 

для i = 0,1,2 и j = 0,1,2.