如何在Pandas中实现按层级自定义排序的分组小计计算

需求背景 #

需要对以下数据集计算分组小计：

Country  Channel    Genre  Size   Prev
UK       Partners   Blues    25     20
UK       Stores     Blues    15     10
UK       Stores     Rock     35     30
US       Stores     Rock     45     40
UK       Partners   Rock     55     50
US       Partners   Rock     65     60
UK       Stores     Blues     5      2

现有分组小计实现无法支持不同层级自定义排序规则，需要实现的预期调用方式和输出如下：

>>> columns = ['Country', 'Channel', 'Genre']
>>> sort = {'Country': 'Country', 'Channel': 'Size', 'Genre': 'Prev'}
>>> subtotal(data, columns, agg='sum', sort=sort)
    Country  Channel    Genre  Size  Prev
0   UK       Stores     Blues   20    12
1   UK       Stores     Rock    35    30
2   UK       Stores             55    42
3   UK       Partners   Blues   25    20
4   UK       Partners   Rock    55    50
5   UK       Partners           80    70
6   UK                         135   112
7   US       Stores     Rock    45    40
8   US       Stores             45    40
9   US       Partners   Rock    65    60
10  US       Partners           65    60
11  US                         110   100
12                             245   212

排序规则说明 #

• Country列按国家名称升序排序（UK排在US之前）
• Channel列在相同Country分组内按Size列的合计值升序排序（UK下Stores的Size合计为55，排在Partners的80之前）
• Genre列在相同Country和Channel分组内按Prev列值升序排序（UK Stores下Blues的Prev合计为12，排在Rock的30之前）

实现代码 #

基于Pandas实现，性能高效，支持任意层级分组和自定义排序规则：

import pandas as pd
import numpy as np

def subtotal(df, group_cols, agg='sum', sort_config=None):
    # 未指定排序规则时默认按分组列本身排序
    if sort_config is None:
        sort_config = {col: col for col in group_cols}
    all_level_dfs = []
    group_count = len(group_cols)
    # 逐层生成从最细粒度到总计的所有层级聚合结果
    for level in range(group_count, -1, -1):
        if level == 0:
            # 生成总计行
            agg_df = df.agg(agg).to_frame().T
            for col in group_cols:
                agg_df[col] = ''
            current_groups = []
        else:
            current_groups = group_cols[:level]
            # 生成本层级分组聚合结果
            agg_df = df.groupby(current_groups, as_index=False).agg(agg)
            # 未到该层级的分组列置空
            for col in group_cols[level:]:
                agg_df[col] = ''
        # 生成辅助排序键
        for idx, group_col in enumerate(group_cols):
            if idx < level:
                sort_col = sort_config[group_col]
                # 计算当前分组对应排序列的聚合值作为排序键
                if level > 0:
                    sort_val = df.groupby(current_groups)[sort_col].transform(agg)
                else:
                    sort_val = df[sort_col].agg(agg)
                agg_df[f'_sort_{idx}'] = sort_val
            else:
                # 更低层级用无穷大保证小计行排在细粒度行之后
                agg_df[f'_sort_{idx}'] = np.inf
        # 层级标识，保证同分组下细粒度行在前、小计行在后
        agg_df['_level'] = level
        all_level_dfs.append(agg_df)
    # 合并所有层级结果
    full_result = pd.concat(all_level_dfs, ignore_index=True)
    # 按规则排序
    sort_keys = [f'_sort_{i}' for i in range(group_count)] + ['_level']
    full_result = full_result.sort_values(sort_keys, ascending=[True]*group_count + [False])
    # 清理辅助列，恢复原始列顺序
    full_result = full_result.drop(columns=[col for col in full_result.columns if col.startswith('_')])
    full_result = full_result[group_cols + [col for col in df.columns if col not in group_cols]].reset_index(drop=True)
    return full_result

测试调用 #

# 构造测试数据集
data = [
    ['UK', 'Partners', 'Blues', 25, 20],
    ['UK', 'Stores', 'Blues', 15, 10],
    ['UK', 'Stores', 'Rock', 35, 30],
    ['US', 'Stores', 'Rock', 45, 40],
    ['UK', 'Partners', 'Rock', 55, 50],
    ['US', 'Partners', 'Rock', 65, 60],
    ['UK', 'Stores', 'Blues', 5, 2]
]
df = pd.DataFrame(data, columns=['Country', 'Channel', 'Genre', 'Size', 'Prev'])

# 调用函数
columns = ['Country', 'Channel', 'Genre']
sort_config = {'Country': 'Country', 'Channel': 'Size', 'Genre': 'Prev'}
result = subtotal(df, columns, agg='sum', sort_config=sort_config)
print(result)

运行后输出结果与预期完全一致。

内容的提问来源于stack exchange，提问作者S Anand