简约的Python索引

Question

我有一个numpy数组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]]]

我想用这样的方法来分割它的一部分：

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]]]

我尝试了方法1

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

和方法2

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

它们的结果都是：

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

有什么方法可以像方法1或2那样处理索引吗？

Answer 1

您可以使用np.all(..., axis=1)

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

Answer 2

您还可以使用列表理解，它允许您适应您的自定义条件。

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