sort a pandas df based on column values that has mix of integer and list of two integers

Viewed 50

I have a dataframe: temp_df

line  sheet  comments
1     apple  this is a fruit
2     orange this is fruit
[1,3] onion  this is a vegetable

I want to sort this temp_df dataframe based on sheet and line but since line has mix of integer and list. I will take the first index if the column value is a list. so the sorted temp_df looks like:

line  sheet  comments
    1     apple  this is a fruit
    [1,3] onion  this is a vegetable
    2     orange this is fruit
1 Answers

You could do:

import pandas as pd

# data setup
data = [[1, 2, [1, 3]], ["apple", "onion", "orange"],
        ["this is a fruit", "this is a vegetable", "this is a fruit"]]
df = pd.DataFrame(data=zip(*data), columns=["line", "sheet", "comments"])

result = df.assign(key=[s if isinstance(s, list) else [s] for s in df["line"]]) \
    .sort_values("key").drop("key", axis=1)

print(result)

Output

     line   sheet             comments
0       1   apple      this is a fruit
2  [1, 3]  orange      this is a fruit
1       2   onion  this is a vegetable

In pandas >= 1.1, you could use the key parameter of sort_values see sorting with keys. Like follows:

result = df.sort_values(by=['line'], key=lambda ix: [s if isinstance(s, list) else [s] for s in ix])
print(result)
Related