从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()获得最大值。这将是一维整数索引，但是您可以使用np.unravel_index将其转换为Pl [5]的二维索引。最后，我们缺少沿已知为5的零位维度的索引。只需在它前面添加并返回数组即可。