Python中的最大值3D数组

我们有一个3D阵列（6719066,2）：

• 671 IdUser（集团）

• 对于每671个用户，我们有9066（行）部电影和相关的proba（2列）

一个用户的两列示例：

Proba                Idmovie
[  6.93933334e-02,   1.00000000e+00],
[  8.50563031e-02,   2.00000000e+00],
....
[  2.95071098e-02,   3.00000000e+00]

我们需要最大概率（此处为8.50563031e-02）和相关第二列的值（此处为2.00000000 E+00）。这适用于每671个用户

---

非常感谢

肯定有更聪明的方法，但我的方法是：

import numpy as np

x = np.random.random((3,3,2))
print(x)


for user in x:
    maximum = np.max(user[:,0])
    index = np.where(user[:,0]==maximum)

    print(user[index])

---

numpy中最紧凑的方式可能是这样的：

x[range(x.shape[0]),np.argmax(x, axis=1)[:,0],:]

说明：

np.argmax(x, axis=1)

返回沿第二（“电影”）轴的最大值的索引

因为我们只需要最大概率的索引，所以将其减少到第一列

使用这些索引作为索引数组，我们只需检索每个用户各自的proba/idmovie元组。结果是一个形状数组（671,2），具有每个用户的最大概率和关联的电影id

np.argmax(x, axis=1)[:,0]