使用Pandas实现DataFrame的压缩与反聚合转换难题

问题背景 #

需要将宽格式DataFrame转换为长格式，并根据季度列的数值反聚合（即按数值生成对应行数的记录），同时为type字段添加序号后缀。

原始数据 #

state   range   type    Q1 24   Q2 24   stat
NY      up      AA      2       2       grow
NY      up      AA      1       0       re
NY      up      BB      1       1       grow
NY      up      BB      0       0       re
NY      up      DD      2       3       grow
NY      up      DD      0       1       re
CA      low     AA      0       2       grow
CA      low     AA      1       0       re
CA      low     BB      0       1       grow
CA      low     BB      0       0       re
CA      low     DD      0       3       grow
CA      low     DD      1       0       re

对应DataFrame构建代码：

import pandas as pd

data = {
    'state': ['NY', 'NY', 'NY', 'NY', 'NY', 'NY', 'CA', 'CA', 'CA', 'CA', 'CA', 'CA'],
    'range': ['up', 'up', 'up', 'up', 'up', 'up', 'low', 'low', 'low', 'low', 'low', 'low'],
    'type': ['AA', 'AA', 'BB', 'BB', 'DD', 'DD', 'AA', 'AA', 'BB', 'BB', 'DD', 'DD'],
    'Q1 24': [2, 1, 1, 0, 2, 0, 0, 1, 0, 0, 0, 1],
    'Q2 24': [2, 0, 1, 0, 3, 1, 2, 0, 1, 0, 3, 0],
    'stat': ['grow', 're', 'grow', 're', 'grow', 're', 'grow', 're', 'grow', 're', 'grow', 're']
}

df = pd.DataFrame(data)

期望输出 #

state   qtr     type    range   stat    
NY      Q1 24   AA01    up      grow    
NY      Q1 24   AA02    up      grow    
NY      Q1 24   AA03    up      re  
NY      Q1 24   BB01    up      grow    
NY      Q1 24   DD01    up      grow    
NY      Q1 24   DD02    up      grow    
CA      Q1 24   AA01    low     re  
CA      Q1 24   DD01    low     re  
NY      Q2 24   AA01    up      grow    
NY      Q2 24   AA02    up      grow    
NY      Q2 24   BB01    up      grow    
NY      Q2 24   DD01    up      grow    
NY      Q2 24   DD02    up      grow    
NY      Q2 24   DD03    up      grow    
NY      Q2 24   DD04    up      re  
CA      Q2 24   AA01    low     grow    
CA      Q2 24   AA02    low     grow    
CA      Q2 24   BB01    low     grow    
CA      Q2 24   DD01    low     grow    
CA      Q2 24   DD02    low     grow    
CA      Q2 24   DD03    low     grow    

完整实现代码 #

以下是一步到位的链式操作，解决列压缩和数值反聚合问题：

import pandas as pd

# 1. 宽表转长表，提取季度和对应计数
df_long = df.pivot_longer(
    cols=['Q1 24', 'Q2 24'],
    names_to='qtr',
    values_to='repeat_count'
)

# 2. 过滤掉不需要展开的0计数行
df_long = df_long[df_long['repeat_count'] > 0]

# 3. 按计数反聚合，重复对应行数
df_exploded = df_long.loc[df_long.index.repeat(df_long['repeat_count'])]

# 4. 为每个分组生成序号，更新type字段
df_exploded['type'] = (
    df_exploded['type'] + 
    df_exploded.groupby(['state', 'qtr', 'type', 'stat', 'range'])
               .cumcount()
               .add(1)
               .astype(str)
               .str.zfill(2)  # 确保序号是两位格式，如01、02
)

# 5. 调整列顺序并重置索引
final_df = df_exploded[['state', 'qtr', 'type', 'range', 'stat']].reset_index(drop=True)

print(final_df)

代码说明 #

1. 宽表转长表：使用pivot_longer将Q1 24、Q2 24列转为qtr（季度）列，对应数值存入repeat_count（需要重复的次数）。
2. 过滤无效行：移除repeat_count为0的行，这些行不需要生成记录。
3. 反聚合行：通过index.repeat按repeat_count的值重复每行，实现数值到行数的转换。
4. 生成带序号的type：按state、qtr、type、stat、range分组，用cumcount生成组内序号，拼接到原type后，用str.zfill(2)保证序号为两位格式。
5. 整理输出：调整列顺序为期望格式，重置索引。

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