NumPy: Как получить индексы максимальных значений в многомерном массиве

Question

Со следующим массивом:

In [103]: da                                                                                         
Out[103]: 
array([[[ 6, 22,  3],
        [ 4,  9, 20],
        [21, 16,  0]],

       [[ 2, 25, 11],
        [ 5, 17, 18],
        [23, 13,  7]],

       [[10, 14, 26],
        [ 8,  1, 19],
        [15, 12, 24]]])

In [104]: da.shape                                                                                   
Out[104]: (3, 3, 3)

Индексы элемента с максимальным значением можно определить с помощью следующего:

In [114]: np.unravel_index(np.argmax(da), da.shape)                                                  
Out[114]: (2, 0, 2)

и проверено:

In [115]: da[2, 0, 2]                                                                                
Out[115]: 26

Но можно ли определить, без зацикливания / повторения, 9 индексов, содержащих максимальные значения для каждой группы целых чисел da[:, i1, i2], где, где i1 и i2 равны 0, 1 или 2?

Например, группа da[:, 0, 0] - это 6, 2 и 10. Максимальное значение - 10, а ее индексы - da[2, 0, 0].

Answer 1

Аргумент axis позволяет указать одну ось операции:

i0 = np.argmax(da, axis=0)

Это означает, что i0 - это массив (3, 3), содержащий индекс максимума для каждой соответствующей i1, i2. Максимальное значение для любого i1, i2 в вашем примере:

da[i0[i1, i2], i1, i2]