Return 一维 numpy isin 函数的掩码

Question

我正在尝试使用 numpy 的函数 isin 来 return 给定查询的掩码。例如，假设我想在下面的 numpy 数组中获取元素 2.1 的掩码：

import numpy as np

a = np.array(
    [
        ["1", "1.1"],
        ["1", "1.2"],
        ["2", "2.1"],
        ["2", "2.2"],
        ["2.1", "2.1.1"],
        ["2.1", "2.1.2"],
        ["2.2", "2.2.1"],
        ["2.2", "2.2.2"],
    ]
)

我正在使用参数 np.isin(a, "2.1") 查询它，但是这个 return 是另一个二维数组而不是一维掩码：

[[False False]
 [False False]
 [False  True]
 [False False]
 [ True False]
 [ True False]
 [False False]
 [False False]]

我原以为它会 return 类似于：

[False False True False True True False False]

我应该怎么做才能解决这个问题？

Answer 1

如果您想要“2.1”出现在 a 中的行，您需要轴上的 any 方法：

>>> np.isin(a, "2.1").any(axis=1)
array([False, False,  True, False,  True,  True, False, False])

如果您想要 a 中出现“2.1”的位置的索引，您可以使用 np.where:

>>> np.where(np.isin(a, "2.1"))
(array([2, 4, 5], dtype=int64), array([1, 0, 0], dtype=int64))

Return 一维 numpy isin 函数的掩码

Return mask from numpy isin function in 1 dimension

python

numpy

isin