Return mask from numpy isin function in 1 dimension-CodePudding

I am trying to use numpy's function isin to return a mask for a given query. For example, let's say I want to get a mask for element 2.1 in the numpy array below:

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"],
    ]
)

I am querying it with the parameters np.isin(a, "2.1"), but this returns another 2D array instead of a 1D mask:

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

I was expecting it would return something like:

[False False True False True True False False]

What am I supposed to do to fix this query?

CodePudding user response：

If you want the rows that "2.1" appears in a, you want the any method on axis:

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

If you want the indexes of where "2.1" appears in a, you could use np.where:

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