Python pandas find element of one column in list of elements of another column

1.9k views Asked by At

Here is my problem, where I want to find element of column A in list of elements of column B of a dataframe. As a result I want to to only keep those rows, where the element in A was found:

df = pd.DataFrame({'A': [1, 2],
                   'B': [1, 3]
                 })
result = df[df.A.isin(df.B)]

>>> result
   A  B
0  1  1

works fine, but what I really want is:

df = pd.DataFrame({'A': [1, 2],
                   'B': [[1, 2], [1, 3]]
                 })
result = df[df.A.isin(df.B)]

>>> result
Empty DataFrame
Columns: [A, B]
Index: []

Which does not work as the elements from A are not compared with the elements of the lists in column B but with the whole list?

What I would like to have as a result is:

>>> result
   A       B
0  1  [1, 2]

Is that possible?

1

There are 1 answers

2
Quang Hoang On BEST ANSWER

You can do apply:

df[df.apply(lambda row: row['A'] in row['B'], axis=1)]

or zip comprehension:

df[[a in b for a,b in zip(df['A'], df['B'])]]

output:

   A       B
0  1  [1, 2]