Python indexing with simplicity-CodePudding

I have a numpy array a,

import numpy as np

a = np.array([[[3, 2, 2], [3, 4, 2]],
[[1, 2, 2], [3, 4, 2]],
[[1, 2, 2], [3, 4, 2]]
])
print(a)

[[[3 2 2]
  [3 4 2]]

 [[1 2 2]
  [3 4 2]]

 [[1 2 2]
  [3 4 2]]]

and I want to slice part of it, right now with this way:

b = []
for i in range(a.shape[0]):
    if (a[i, 0, 0] > 2 and a[i, 1, 0] > 2):
        b.append(a[i])
print(np.array(b))

[[[3 2 2]
  [3 4 2]]]

I tried method 1

a[np.where(a[:,:,0] > 2)]

and method 2

a[a[:,:,0]> 2]

They both result in:

array([[3, 2, 2],
       [3, 4, 2],
       [3, 4, 2],
       [3, 4, 2]])

Is there any way to deal with the index like method 1 or 2?

CodePudding user response：

You can use np.all(..., axis=1):

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

CodePudding user response：

You could also use a list comprehension which allow you to fit your custom conditions.

b = [i for i in a if i[0,0]>2 and i[1,0]>2]