根据行元素是否在另一个数组中过滤numpy数组的行

Question

我有一个数组 group，它是 Nx2：

array([[    1,     6],
       [    1,     0],
       [    2,     1],
       ...,
       [40196, 40197],
       [40196, 40198],
       [40196, 40199]], dtype=uint32)

还有另一个数组selection，即(M,)：

array([3216, 3217, 3218, ..., 8039]) 

我想创建一个包含group 的所有行的新数组，其中两个元素都在selection 中。我就是这样做的：

np.array([(i,j) for (i,j) in group if i in selection and j in selection])

这可行，但我知道必须有一种更有效的方法来利用一些 numpy 函数。

Answer 1

您可以使用np.isin 获得一个与group 形状相同的布尔数组，该数组表示元素是否在selection 中。然后，要检查行中的两个条目是否都在selection 中，您可以使用all 和axis=1，这将给出一个一维布尔数组，说明要保留哪些行。我们终于用它索引了：

group[np.isin(group, selection).all(axis=1)]

---

示例：

>>> group

array([[    1,     6],
       [    1,     0],
       [    2,     1],
       [40196, 40197],
       [40196, 40198],
       [40196, 40199]])

>>> selection

array([    1,     2,     3,     4,     5,     6, 40196, 40199])

>>> np.isin(group, selection)

array([[ True,  True],
       [ True, False],
       [ True,  True],
       [ True, False],
       [ True, False],
       [ True,  True]])

>>> np.isin(group, selection).all(axis=1)

array([ True, False,  True, False, False,  True])

>>> group[np.isin(group, selection).all(axis=1)]

array([[    1,     6],
       [    2,     1],
       [40196, 40199]])