python - 从一维的numpy isin函数返回掩码

Question

我正在尝试使用 numpy 的函数 isin 返回给定查询的掩码。例如，假设我想在下面的 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")，但这会返回另一个 2D 数组而不是 1D 掩码：

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

我期待它会返回类似的东西：

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

如果你想要“2.1”出现在哪里的索引a，你可以使用np.where：

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