如何在 python 中按条件使用 select 特定索引

Question

我有以下包含 3 个子数组的数组：

array([[1,3],[2,3],[4,5]],
      [[4],[7],[9]],
      [[0.3],[0.02],[-0.001]])

我想创建一个新数组，其中仅包含符合以下条件的索引（在每个子数组中）： abs(array[2]) < 0.1 - 这意味着我最终希望收到 output：

array([[2,3],[4,5]],
      [[7],[9]],
      [[0.02],[-0.001]])

有什么建议可以有效地做到这一点吗？

Answer 1

我会这样做：

import numpy as np

a = np.array([[[1,3],[2,3],[4,5]],
      [[4],[7],[9]],
      [[0.3],[0.02],[-0.001]]])

b = [abs(item[0])<0.1 for item in a[2]]

c = [row[b] for row in a]

b 是一个布尔数组，定义指向 select 的索引