时间戳

首页 > 解决方案 > 根据行元素是否在另一个数组中过滤numpy数组的行

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

问题描述

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 函数。

标签： pythonarraysperformancenumpy

解决方案

您可以使用它来获取与表示元素是否在 中np.isin的形状相同的布尔数组。然后，要检查行中的两个条目是否都在 中，您可以使用with ，这将给出一个一维布尔数组，说明要保留哪些行。我们最终用它索引：groupselectionselectionallaxis=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]])

推荐阅读