如果满足条件，如何将数组的值设置为数字？

Question

我的矩阵A的大小为2x4x3x3，其中2表示两个矩阵A1和A2的组合。对于每个矩阵A1/A2（大小为4x3x3），我只选择大于0.3（小于0.3的值，称为不确定性位置）并执行argmax以获得矩阵C。在矩阵C中，我将不确定性位置上的值设置为255。我在python中通过以下两种方法进行了设置。我正在寻找第三种更短的方法来将不确定性位置的值设置为255。可以吗？谢谢

代码是

import numpy as np
A1=np.array([[[0.4,0.2,0.3],
            [0.4,0.5,0.6],
            [0.7,0.8,0.2]],
            [[0.4,0.5,.3],
            [0.2,0.5,0.2],
            [0.3,0.2,0.1]],
            [[0.5,0.1,0.6],
            [0.8,0.1,0.5],
            [0.5,0.1,0.4]],
            [[0.6,0.1,0.1],
            [0.2,0.1,0.9],
            [0.9,0.1,0.1]]])
A2=np.array([[[0.8,0.1,0.6],
            [0.4,0.6,0.6],
            [0.7,0.8,0.2]],
            [[0.4,0.8,.3],
            [0.2,0.8,0.2],
            [0.3,0.0,0.1]],
            [[0.5,0.1,0.6],
            [0.8,0.0,0.5],
            [0.5,0.1,0.4]],
            [[0.6,0.1,0.1],
            [0.2,0.1,0.1],
            [0.9,0.1,0.1]]])

A1= np.expand_dims(A1,0)
A2= np.expand_dims(A2,0)
A= np.concatenate([A1,A2],axis=0)  
B = A< 0.3

#First way
C=np.argmax(A, axis=1)
C[B[:,0,:,:]]=255
C[B[:,1,:,:]]=255
C[B[:,2,:,:]]=255
C[B[:,3,:,:]]=255
print (C)

#Second way
C=np.argmax(A, axis=1)
for i in range (4):
    C[B[:,i,:,:]] = 255
print (C)

#Third way
C=np.argmax(A, axis=1)
C[B[:,0:3,:,:]]=255
print (C)

您也可以运行它Online Code