Extract the first 2 dimention of numpy array based on condtion on the 3rd dimention-CodePudding

Suppose that I have numpy array with 3 dimension [x,y,z], and I would like to extract the x and y dimension on a condition on the z dimension, for example, if z==1. How Could I do that?

CodePudding user response：

You can use numpy indexing for this. Consider an array:

a = np.arange(21).reshape((-1, 3))

# array([[ 0,  1,  2],
#       [ 3,  4,  5],
#       [ 6,  7,  8],
#       [ 9, 10, 11],
#       [12, 13, 14],
#       [15, 16, 17],
#       [18, 19, 20]])

Now you want a condition on the last column...say ever numbers:

    # all rows ⬎  ⬐ third column  
filtered = a[a[:, 2] % 2 == 0]
# array([[ 0,  1,  2],
#       [ 6,  7,  8],
#       [12, 13, 14],
#       [18, 19, 20]])

And just select the first two columns:

filtered[:,:2]   
# array([[ 0,  1],
#        [ 6,  7],
#        [12, 13],
#        [18, 19]])

This works because this gives an array of booleans...

i = a[:, 2] % 2 == 0
# array([ True, False,  True, False,  True, False,  True])

...which can then be used to index the original:

a[i]
# array([[ 0,  1,  2],
#        [ 6,  7,  8],
#        [12, 13, 14],
#        [18, 19, 20]])