Pandas：避免groupby-apply操作后结果按索引排序

问题背景 #

原始DataFrame定义如下：

import pandas as pd
import math

dict_ = {
    'Query' : ['apple', 'banana', 'mango', 'bat', 'cat', 'rat', 'lion', 'potato', 'london', 'new jersey'],
    'Category': ['fruits', 'fruits', 'fruits', 'animal', 'animal', 'animal', 'animal', 'veggie', 'place', 'place'],
}
df = pd.DataFrame(dict_)

需求是按Category分组后复制每组数据，让所有组的最终行数等于最大组的行数（此处最大组为animal，共4行）。但执行以下代码后，结果的分组顺序变为animal、fruits、place、veggie，不符合原DataFrame中分组出现的顺序fruits、animal、veggie、place：

rep_val= df.groupby('Category').size().max()
df.groupby('Category').apply(lambda d: pd.concat(([d]*math.ceil(rep_val/d.shape[0]))).head(rep_val)).reset_index(drop=True)

预期输出：

Query Category
0        apple   fruits
1       banana   fruits
2        mango   fruits
3        apple   fruits
4          bat   animal
5          cat   animal
6          rat   animal
7         lion   animal
8       potato   veggie
9       potato   veggie
10      potato   veggie
11      potato   veggie
12      london    place
13  new jersey    place
14      london    place
15  new jersey    place

解决方法 #

Pandas默认groupby会对分组键排序，要保留原始出现顺序，只需在groupby时添加sort=False参数。修改后的代码如下：

import pandas as pd
import math

dict_ = {
    'Query' : ['apple', 'banana', 'mango', 'bat', 'cat', 'rat', 'lion', 'potato', 'london', 'new jersey'],
    'Category': ['fruits', 'fruits', 'fruits', 'animal', 'animal', 'animal', 'animal', 'veggie', 'place', 'place'],
}
df = pd.DataFrame(dict_)

rep_val = df.groupby('Category', sort=False).size().max()
result = df.groupby('Category', sort=False).apply(
    lambda d: pd.concat([d]*math.ceil(rep_val/d.shape[0])).head(rep_val)
).reset_index(drop=True)

print(result)

关键说明 #

• 两次调用groupby时都添加sort=False：第一次计算每组大小最大值时，保证按原始顺序统计；第二次分组处理数据时，保证分组顺序与原DataFrame中Category首次出现的顺序一致。
• 执行修改后的代码，输出结果的分组顺序将与预期完全匹配。

内容的提问来源于stack exchange，提问作者Himanshu Poddar