Clearing a row based on the state of the first element

Viewed 39

There is a 3d array:

Input:

    [[[0,2,3,4]
    [4,2,3,4]
    [6,2,3,4]]
    [[2,2,3,4]
    [3,2,3,4]
    [2,2,3,4]]]

How can I make a numpy array look like this?

rule: if array[:,:,0] < 3

Output:

    [[[0,0,0,0]
    [4,2,3,4]
    [6,2,3,4]]
    [[0,0,0,0]
    [0,0,0,0]
    [0,0,0,0]]]
1 Answers

Here's one way:

a[a[:,: ,0] <= 3, :] = 0
OUTPUT:
array([[[0, 0, 0, 0],
        [4, 2, 3, 4],
        [6, 2, 3, 4]],

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

NOTE: I've assumed you wanna turn rows to 0 where the value is < or = to 3. Change the condition if required.

Related