python - numpy 3D 数组中沿平面的最大值索引

Question

我有一个形状为 (3,3,3) 的 3D numpy 数组。我想获得平面中最大值的索引，根据我的说法，“平面”如下:

a = np.random.rand(3,3,3)
>>> a[:,:,0]
array([[0.98423332, 0.44410844, 0.06945133],
       [0.69876575, 0.87411547, 0.53595041],
       [0.53418486, 0.16186808, 0.60579623]])
>>> a[:,:,1]
array([[0.38969199, 0.80202126, 0.62189662],
       [0.66609605, 0.09771614, 0.74061269],
       [0.77081531, 0.20068743, 0.72762023]])
>>> a[:,:,2]
array([[0.57110332, 0.29021439, 0.15433043],
       [0.21762439, 0.93112448, 0.05763075],
       [0.77880124, 0.36637245, 0.29070822]])

我有一个解决方案，但我想要更短、更快、没有 for 循环的解决方案，我的解决方案如下:

for i in range(3):
    x=a[:,:,i].argmax()/3
    y=a[:,:,i].argmax()%3
    z=i
    print(x,y,z)
    print a[x][y][z]
(0, 0, 0)
0.9842333247061394
(0, 1, 1)
0.8020212566990867
(1, 1, 2)
0.9311244845473187

Answer 1

我们只需要将输入数组 reshape 为2D，方法是合并最后两个轴，然后沿第二个轴(即合并后的轴)应用argmax，为我们自己提供一个矢量化的数组方法-

def argmax_each_plane(a):
    a2D = a.reshape(a.shape[0],-1)
    idx = a2D.argmax(1)
    indices = np.unravel_index(idx, a.shape[1:])
    vals = a2D[np.arange(len(idx)), idx]
    return vals, np.c_[indices]

示例运行 -

In [60]: np.random.seed(0)
    ...: a = np.random.rand(3,3,3)

In [61]: a
Out[61]: 
array([[[0.5488135 , 0.71518937, 0.60276338],
        [0.54488318, 0.4236548 , 0.64589411],
        [0.43758721, 0.891773  , 0.96366276]],

       [[0.38344152, 0.79172504, 0.52889492],
        [0.56804456, 0.92559664, 0.07103606],
        [0.0871293 , 0.0202184 , 0.83261985]],

       [[0.77815675, 0.87001215, 0.97861834],
        [0.79915856, 0.46147936, 0.78052918],
        [0.11827443, 0.63992102, 0.14335329]]])

In [62]: v, ind = argmax_each_plane(a)

In [63]: v
Out[63]: array([0.96366276, 0.92559664, 0.97861834])

In [64]: ind
Out[64]: 
array([[2, 2],
       [1, 1],
       [0, 2]])

如果您还需要 z 索引，请使用:np.c_[indices[0],indexs[1], range(len(a2D))]。