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

我们只需要通过合并最后两个轴，然后沿第二个轴应用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], indices[1], range(len(a2D))]。

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

1 个答案: