CODEWITHSUNDEEP
pythonlistpandas
NBC
NBC

Reputation: 1708

Comparisons with Pandas DataFrame columns of lists

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

Upvotes: 3

Views: 118

Answers (1)

cs95
cs95

Reputation: 403128

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.

Upvotes: 4

Related Questions