Заменить весь массив, если значение больше

Question

У меня есть массив ниже с формой (2, 6, 3)

[[[ 10.          0.          0.        ]
  [ 20.          0.          0.        ]
  [ 60.          0.          0.        ]
  [ 40.          0.          0.        ]
  [ 20.          0.          0.        ]
  [ 80.          0.          0.        ]]

 [[ 70.          0.          0.        ]
  [ 90.          0.          0.        ]
  [ 20.          0.          0.        ]
  [ 50.          0.          0.        ]
  [ 80.          0.          0.        ]
  [ 10.          0.          0.        ]]]

Я пытаюсь заменить полный массив, если 1-е значение> = 50 с [1 1 1], и еслипервое значение <50, тогда весь массив должен иметь значение [0 0 0] </p>

Желаемый вывод:

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

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

Я могу заменить только одно значение, но не полное.

Любая идея о том, как это можно сделать

Answer 1

Вот один подход индексирования из массива, содержащего [0,0,0] и [1,1,1] на основе результата сравнения:

b = np.array([[0,0,0], [1,1,1]])

b[(a[:,:,0] >= 50).view('i1')]

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

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