NumPy: как получить максимум от Argmax результате

Question

У меня есть Numpy массив произвольной формы, например:

a = array([[[ 1, 2], 
      [ 3, 4], 
      [ 8, 6]], 

      [[ 7, 8], 
      [ 9, 8], 
      [ 3, 12]]]) 
a.shape = (2, 3, 2) 

и результат Argmax над последней осью:

np.argmax(a, axis=-1) = array([[1, 1, 0], 
           [1, 0, 1]]) 

I Желательно получить максимум:

np.max(a, axis=-1) = array([[ 2, 4, 8], 
          [ 8, 9, 12]]) 

Но без пересчета все. Я пробовал:

a[np.arange(len(a)), np.argmax(a, axis=-1)] 

Но получил:

IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (2,) (2,3) 

Как это сделать? Похожий вопрос на 2-й: numpy 2d array max/argmax

Answer 1

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

In [17]: a 
Out[17]: 
array([[[ 1, 2], 
     [ 3, 4], 
     [ 8, 6]], 

     [[ 7, 8], 
     [ 9, 8], 
     [ 3, 12]]]) 

In [18]: idx = a.argmax(axis=-1) 

In [19]: m,n = a.shape[:2] 

In [20]: a[np.arange(m)[:,None],np.arange(n),idx] 
Out[20]: 
array([[ 2, 4, 8], 
     [ 8, 9, 12]]) 

Для общего ndarray случае любого числа измерений, как указано в comments by @hpaulj, мы могли бы использовать np.ix_, как так -

shp = np.array(a.shape) 
dim_idx = list(np.ix_(*[np.arange(i) for i in shp[:-1]])) 
dim_idx.append(idx) 
out = a[dim_idx]