Как установить значения массива в число, если оно удовлетворяет условию?

Question

У меня есть матрица размером 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)

Вы также можете запустить его Код онлайн