How to check inner part of 3D array numpy-CodePudding

I have a 3D array a = np.array([[[1,1],[1,1]],[[1,0],[1,0]]]) and I need to check which members are equal to [1,1] and set all other members to [0,0] How should I approach this?

CodePudding user response：

You can use numpy.where with the appropriate axis to find the subarrays. Given that you can assign a new value:

import numpy as np

a = np.array([[[1, 1],
            [1, 1]],

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

           [[1, 1],
            [1, 0]]
])
a[np.where(np.all(a == [1,1], axis=2))] = [0, 0]

This will give you an a with:

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

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

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