与Pandas DataFrame列列表进行比较

Question

I have a dataframe df like this:

col1 | col2
 a   | [1,2]
 b   | [3,4]
 c   | [3,9]

I want to get the row based on a matching input array, so if I have the array [1,2], I can get:

col1 | col2
 a   | [1,2]

When I try to do this using this formula, it doesn't work:

df.loc[df['Col2'] == [1,2]]
Error: Lengths must match to compare

Answer 1

The real cause of your error was that not all lists are of the same size, this causes issues with DataFrame.eq .

The best approach to tackling this is to build a boolean mask using a list comprehension and then use it to index into df :

df[[v == [1, 2] for v in df['col2'].tolist()]]

Another alternative is df.apply , but that isn't nearly as fast as this.