Выберите строки с самыми высокими значениями из np.zeros

Question

У меня есть следующая структура данных:

Pl = np.zeros((7,2,7))
Pl[0,0,1]=1
Pl[1,0,2]=1
Pl[2,0,3]=1
Pl[5,0,6]=0.9
Pl[5,0,5]=0.1
...
Pl[5,1,4]=1

Как я могу получить запись с указанным первым значением, которое имеет самое высокое присвоенное значение?

Например, для x = 5 я хочу получить Pl [5,1,4].Я видел max, но не могу указать значение x.

Спасибо!

Answer 1

import numpy as np

Pl = np.zeros((7, 2, 7))
Pl[0, 0, 1] = 1
Pl[1, 0, 2] = 1
Pl[2, 0, 3] = 1
Pl[5, 0, 6] = 0.9
Pl[5, 0, 5] = 0.1
Pl[5, 1, 4] = 1

res = np.array([5, *np.unravel_index(Pl[5].argmax(), Pl[5].shape)])
# array([5, 1, 4])

Здесь Pl[5].argmax() получает максимальное значение.Это будет одномерный целочисленный индекс, но вы можете преобразовать его в двумерный индекс Pl [5] с помощью np.unravel_index.Наконец, нам не хватает индекса по нулевому измерению, которое, как мы знаем, равно 5.Просто добавьте его и верните массив.