如何用pandas原生方法生成指定格式的统计DataFrame？

原始数据 #

| code | company | user | all | cat_1 | cat_2 | cat_3 | price |
| ---- | ------- | ---- | --- | ----- | ----- | ----- | ----- |
| ABC  | 1       | 123  | x   | x     |       | x     | 50    |
| ABC  | 1       | 456  | x   |       | x     |       | 70    |
| ABC  | 1       | 789  | x   | x     | x     |       | 90    |
| ABC  | 2       | 098  | x   |       |       | x     | 55    |
| ABC  | 2       | 765  | x   | x     |       |       | 75    |
| ABC  | 2       | 432  | x   | x     | x     | x     | 95    |

期望输出 #

| code | company | cat   | price_n | price_avg |
| ---- | ------- | ----- | ------- | --------- |
| ABC  | 1       | all   | 3       | 70        |
| ABC  | 1       | cat_1 | 2       | 70        |
| ABC  | 1       | cat_2 | 2       | 80        |
| ABC  | 1       | cat_3 | 1       | 50        |
| ABC  | 2       | all   | 3       | 75        |
| ABC  | 2       | cat_1 | 2       | 85        |
| ABC  | 2       | cat_2 | 1       | 95        |
| ABC  | 2       | cat_3 | 2       | 75        |

现有实现思路 #

目前通过循环遍历['all', 'cat_1', 'cat_2', 'cat_3']字段，对每个字段分组聚合后合并结果（类似SQL的UNION操作），补全后的代码如下：

import pandas as pd

def crunch_data(df: pd.DataFrame) -> pd.DataFrame:
    dfs = []
    for var in ['all', 'cat_1', 'cat_2', 'cat_3']:
        temp = df.groupby(['code', 'company', var]).agg(
            price_n=('user', 'count'),
            price_avg=('price', 'mean')
        ).reset_index()
        temp.rename(columns={var: 'cat'}, inplace=True)
        temp = temp[temp['cat'] == 'x']
        dfs.append(temp)
    result = pd.concat(dfs).sort_values(['code', 'company', 'cat']).reset_index(drop=True)
    result['price_avg'] = result['price_avg'].astype(int)
    return result

更优的原生方法：使用melt重塑数据 #

无需循环，利用Pandas的melt函数将类别列转为行结构，再一次性完成聚合，代码更简洁高效：

import pandas as pd

def crunch_data_optimized(df: pd.DataFrame) -> pd.DataFrame:
    # 1. 重塑数据：把横向类别列转为纵向，保留核心字段作为标识
    melted = df.melt(
        id_vars=['code', 'company', 'user', 'price'],
        value_vars=['all', 'cat_1', 'cat_2', 'cat_3'],
        var_name='cat',
        value_name='flag'
    )
    
    # 2. 过滤仅含有效标记（x）的行
    filtered = melted[melted['flag'] == 'x']
    
    # 3. 分组聚合计算统计值
    result = filtered.groupby(['code', 'company', 'cat']).agg(
        price_n=('user', 'count'),
        price_avg=('price', 'mean')
    ).reset_index()
    
    # 4. 调整格式，匹配示例输出的整数平均值
    result['price_avg'] = result['price_avg'].astype(int)
    
    # 5. 按要求排序
    result = result.sort_values(['code', 'company', 'cat']).reset_index(drop=True)
    
    return result

方法说明 #

• melt：一次性完成多列转置，替代循环拼接的UNION逻辑，代码更简洁且性能更优（数据量大时优势明显）。
• 过滤步骤：仅保留有有效标记的行，避免无效数据干扰统计结果。
• 单次聚合：分组后直接计算所需统计值，无需多次分组操作。

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