如何删除行数超阈值的分组并重塑Pandas DataFrame结构

步骤概述 #

1. 按指定多字段分组，过滤掉行数超过阈值的分组
2. 为每个保留分组内的行添加序号，用于转宽表的列标识
3. 将长表数据重塑为宽表格式

示例代码（基于Pandas） #

假设你的数据集结构如下（可替换为实际数据）：

import pandas as pd

# 示例输入数据
df = pd.DataFrame({
    'ID': ['A', 'A', 'A', 'A', 'A', 'B', 'B', 'B', 'C', 'C', 'C', 'C'],
    'X': [1, 1, 1, 1, 1, 2, 2, 2, 3, 3, 3, 3],
    'Y': ['X', 'X', 'X', 'X', 'X', 'Y', 'Y', 'Y', 'Z', 'Z', 'Z', 'Z'],
    'Z': [10, 10, 10, 10, 10, 20, 20, 20, 30, 30, 30, 30],
    'Value': [100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 1100, 1200]
})

1. 过滤分组（保留行数≤4的分组）

# 按ID、X、Y、Z分组，过滤出行数不超过4的分组
threshold = 4
filtered_df = df.groupby(['ID', 'X', 'Y', 'Z']).filter(lambda group: len(group) <= threshold)

• groupby.filter()会保留所有符合条件的分组的行，适合大型数据集，性能优于先统计分组大小再合并的方式。

2. 添加组内序号

# 为每个分组内的行生成从1开始的序号，用于转宽表时区分不同行的Value
filtered_df['seq'] = filtered_df.groupby(['ID', 'X', 'Y', 'Z']).cumcount() + 1

3. 重塑为宽表

# 转宽表：分组字段作为索引，序号作为列，Value作为单元格值
wide_df = filtered_df.pivot(
    index=['ID', 'X', 'Y', 'Z'],
    columns='seq',
    values='Value'
).reset_index()

# 重命名列，让结果更易读（比如Value_1、Value_2...）
wide_df.columns = ['ID', 'X', 'Y', 'Z'] + [f'Value_{i}' for i in wide_df.columns[4:]]

预期结果 #

执行后wide_df的输出如下：

性能优化建议（针对大型数据集） #

• 如果数据量极大，可先通过groupby.size()统计分组大小，再筛选出符合条件的分组键，最后通过merge过滤原数据，避免lambda可能带来的性能损耗：

  # 先统计每个分组的行数
  group_sizes = df.groupby(['ID', 'X', 'Y', 'Z']).size().reset_index(name='count')
  # 筛选出符合条件的分组键
  valid_groups = group_sizes[group_sizes['count'] <= threshold][['ID', 'X', 'Y', 'Z']]
  # 过滤原数据
  filtered_df = pd.merge(df, valid_groups, on=['ID', 'X', 'Y', 'Z'])
  

• 转宽表时，若存在重复的组内序号（实际不会出现，因为cumcount是唯一的），可使用pivot_table并指定aggfunc='first'确保结果正确。

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