How can I extract rows in Dataframe with function?-CodePudding

My dataframe includes list, like this.

       a        b
1     frog   [1, 2, 3]
2      dog   [4, 5]
3    melon   [6, 7, 1]

I want to extract rows which b contains specific numbers, so I made this function.

def a(_list, _tag):
    if _tag in _list:
        return True
    else:
        return False

I tried to use df.loc[], but it doesn't work well. How can I write a code without iterating all of dataframe?

CodePudding user response：

Here's a way to do what your question asks:

target = 1
df2 = df.explode('b').b == target
df['found'] = df2.groupby(df2.index).sum() > 0
print(df)

Output:

       a          b  found
0   frog  [1, 2, 3]   True
1    dog     [4, 5]  False
2  melon  [6, 7, 1]   True

CodePudding user response：

Use .explode in conjunction with .loc to get relevant indeces, and then query them with .iloc:

df.iloc[df.explode('b').loc[df['b'].isin(target_hash_set)].index.unique()]