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