如何在Python数据框中创建数字后缀列并复现SAS相关功能？

需求说明 #

按Name分组，将每组的Grade值转换为带两位数字后缀的列（最多保留4列），每组仅输出一条汇总记录。

原始数据集 #

Name     Grade
Lee      A
Lee      A+
Lee      A+
Lee      A+
Col      B+
Col      A
Col      B+

目标结果 #

Name    _Grade01    _Grade02    _Grade03    _Grade04
Lee     A           A+          A+          A+
Col     B+          A           B+

对应SAS实现代码 #

array _Grade{*} _Grade01-_Grade04;
if first.Name then do;
    n = 0;
    do _i = 1 to 4;
        _Grade[_i] = 0;
    end;
end;

n + 1;

if n>0 and n <=4 then do;
    _Grade[n] = Grade;
end;

if last.Name then do;
    output;
end;

Python实现方案 #

使用Pandas可以高效实现对应逻辑，步骤如下：

1. 导入库并构造原始数据框

import pandas as pd

data = {
    'Name': ['Lee', 'Lee', 'Lee', 'Lee', 'Col', 'Col', 'Col'],
    'Grade': ['A', 'A+', 'A+', 'A+', 'B+', 'A', 'B+']
}
df = pd.DataFrame(data)

2. 分组添加序号并过滤前4条记录

# 按Name分组，生成组内从1开始的序号
df['seq'] = df.groupby('Name').cumcount() + 1
# 仅保留每组前4条数据（对应SAS中n<=4的判断）
df = df[df['seq'] <= 4]

3. 透视转宽并整理列名

# 透视将序号转为列，Grade值填充对应位置
result = df.pivot(index='Name', columns='seq', values='Grade')
# 重命名列为_GradeXX格式
result.columns = [f'_Grade{str(col).zfill(2)}' for col in result.columns]
# 重置索引，将Name从索引转为普通列
result = result.reset_index()
# 填充缺失值（可根据需求替换为空字符串或其他默认值）
result = result.fillna('')

4. 查看结果

print(result)

逻辑对应说明 #

• groupby('Name').cumcount() +1 对应SAS中first.Name重置n并递增的逻辑，自动为每组生成从1开始的序号
• 过滤seq<=4对应SAS中仅处理前4条记录的判断
• 透视操作对应SAS中数组赋值_Grade[n] = Grade，将组内Grade值映射到对应序号列
• 最终结果自动按组聚合为单条记录，对应SAS中last.Name时输出的逻辑

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