如何用Python编写自定义函数实现DataFrame的排序与透视表生成

你原来的代码存在几个可优化的问题：

1. 函数内部没有使用传入的参数data，反而直接引用了全局变量df，不符合函数封装的要求
2. 列选择语法不合法，Pandas不支持在列选择列表里混用固定列名和列切片'Income_Yr1':'Income_Yr3'
3. 透视表逻辑重复冗余，重复调用了9次pivot_table，且同时将Location放在index和columns参数里，结果结构会出现不必要的冗余
4. 所有透视表生成后没有返回或者打印，调用函数不会输出任何结果

下面是优化后的完整实现：

import pandas as pd
import numpy as np

def process_income_df(data):
    # 1. 选中目标列
    income_cols = [col for col in data.columns if col.startswith('Income_Yr')]
    target_cols = ['Location', 'Group'] + income_cols
    selected_df = data[target_cols].copy()

    # 2. 按Group列升序排序
    sorted_df = selected_df.sort_values(by='Group', ascending=True)

    # 3. 生成合并统计透视表，一次性计算所有年份三个指标的统计值
    result_pivot = pd.pivot_table(
        sorted_df,
        index=['Group', 'Location'],
        values=income_cols,
        aggfunc=['mean', 'median', np.std],
        margins=True
    )

    return sorted_df, result_pivot

# 测试代码
if __name__ == '__main__':
    # 初始化测试数据集
    data = {'Gender':['F', 'F', 'M', 'F','M', 'F', 'M', 'M','F', 'F', 'M', 'F','M', 'F', 'M', 'M','M','F', 'F', 'M'],
        'UID':[1001,1002,1003,1004,1005,1006,1007,1008,1009,1010,1011,1012,1013,1014,1015,1016,1017,1018,1019,1020],
        'Location':['PHX','PHX','PHX','PHX','ATL','ATL','ATL','ATL','HOU','HOU','HOU','MIA','MIA','MIA','MIA','MIA','DEN','DEN','DEN','DEN'],
         'Group':[3,3,3,3,4,4,4,4,1,1,1,1,2,2,2,2,5,5,5,5],
        'Income_Yr1':[32112,34214,45575,22106,32612,34216,47515,22906,32112,34511,45525,12106,52112,54214,45015,22986,32112,34214,47518,22175],
         'Income_Yr2':[52112,54215,65515,72109,52616,64217,77515,52906,52145,38512,65516,32157,63152,57218,51017,42997,38125,36253,49589,32598],
        'Income_Yr3':[52143,54239,65557,72116,52660,64273,77551,52969,52500,38201,65169,32795,63288,57180,51173,42970,38205,36301,59591,32580]}
    df = pd.DataFrame(data)
    sorted_data, pivot_res = process_income_df(df)
    # 输出透视表
    print(pivot_res)

如果你需要单独输出每个指标的透视表，只需要对返回的pivot_res按列筛选即可，比如筛选均值数据：pivot_res['mean']，筛选第一年收入的中位数：pivot_res['median', 'Income_Yr1']。

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