python - 将 3D NumPy 数组中沿轴的最大值设置为 1，其余设置为零

Question

我有一个 3D 数组:

volts = np.random.random((3,3,3)).round(decimals=5)
>>> volts
array([[[0.94785, 0.43955, 0.74527],
    [0.82098, 0.52509, 0.67954],
    [0.72355, 0.16252, 0.03184]],

   [[0.25782, 0.04191, 0.6689 ],
    [0.18215, 0.63108, 0.52052],
    [0.81992, 0.36301, 0.66629]],

   [[0.90585, 0.27223, 0.78807],
    [0.32251, 0.65861, 0.70398],
    [0.21687, 0.20798, 0.33868]]])

对于我的应用程序来说，5 位小数就足够了。

>>> volts[0,0,:]
array([0.94785, 0.43955, 0.74527])  
>>> volts[0,1,:]
array([0.82098, 0.52509, 0.67954])

在上面两行中，我想将 0.94785 和 0.82098 设置为 1，并将所有元素设置为零。不仅是 volts[0,0,:] 和 volts[0,1,:]，还有所有其他 volts[x,y,:]。所以我这样做了:

>>> volts = np.random.random((3,3,3)).round(decimals=5)
>>> volts1 = deepcopy(volts)
>>> vmaxs=volts1.max(axis=2).flatten().tolist()
>>> for items in vmaxs:
       volts1[np.where(volts1==items)]=1


>>> volts1[np.where(volts1!=1)]=0
>>> volts1
array([[[1., 0., 0.],
        [0., 0., 1.],
        [0., 1., 0.]],

       [[0., 1., 0.],
        [1., 0., 0.],
        [1., 0., 0.]],

       [[1., 0., 0.],
        [0., 0., 1.],
        [0., 1., 0.]]])
>>> volts[0,0,:]
array([0.90763, 0.38579, 0.25768])
>>> volts1[0,0,:]
array([1., 0., 0.])
.
.
.
>>> volts[2,2,:]
array([0.33343, 0.73859, 0.43735])
>>> volts1[2,2,:]
array([0., 1., 0.])

您可以看到沿轴 2 的最大值设置为 1，其余元素设置为零。在这里，我只迭代了 9 个元素，但如果我必须迭代 200-300 个元素怎么办？如何以更简洁、更高效的方式做到这一点？

Answer 1

使用np.eye。在有多个最大值的情况下，这将选择第一个。

np.eye(volts.shape[1])[volts.argmax(2)]

array([[[1., 0., 0.],
        [1., 0., 0.],
        [1., 0., 0.]],

       [[0., 0., 1.],
        [0., 1., 0.],
        [1., 0., 0.]],

       [[1., 0., 0.],
        [0., 0., 1.],
        [0., 0., 1.]]])