基于正则聚合Pandas DataFrame行：解决重复ID的数据清洗问题

问题描述 #

我的Pandas DataFrame索引格式为「ID: 名称」，但名称存在输入不规范问题（大小写不一致、名称内容差异），导致同一ID对应多行数据。需要完成两个操作：

• 对同一ID下的所有行数值列求和
• 为每个ID保留任意一个有效名称（如ID=100保留Test或test均可）

环境 #

• Python 3.10.5
• Pandas 1.4.3

示例代码与当前状态 #

初始化代码 #

import pandas as pd

data = [[1, 2], [2, 4], [3, 6], [4, 8], [5, 10]]
index = ['100: Test', '100: test', '101: FOO', '102: WWW', '101: foo foo']
columns = ['column1', 'column2']
df = pd.DataFrame(data, index=index, columns=columns)

print(df)

当前输出 #

column1  column2
100: Test           1        2
100: test           2        4
101: FOO            3        6
102: WWW            4        8
101: foo foo        5       10

期望输出 #

column1  column2
100: Test        3        6  # 注：原示例期望输出数值有误，正确求和应为1+2=3，2+4=6
101: FOO         8       16
102: WWW         4        8

尝试的失败代码 #

直接按原索引分组无法合并同一ID的行：

df2 = df.groupby(level=0).sum()
print(df2)

失败输出 #

column1  column2
100: Test           1        2
100: test           2        4
101: FOO            3        6
101: foo foo        5       10
102: WWW            4        8

解决方案 #

核心是从索引中提取统一的ID作为分组键，同时保留每个ID对应的任意一个原始名称作为新索引。

方法一：正则提取ID分组 #

# 用正则从索引中提取冒号前的ID部分
id_keys = df.index.str.extract(r'^(\d+):', expand=False)

# 按ID分组求和，同时保留每个组的第一个原始索引作为新索引
result_df = df.groupby(id_keys).agg(
    column1=('column1', 'sum'),
    column2=('column2', 'sum'),
    original_index=('column1', lambda x: df.index[x.index[0]])
).set_index('original_index').rename_axis(None)

print(result_df)

方法二：分割索引提取ID分组（更简洁） #

# 按索引中冒号前的ID分组求和
grouped = df.groupby(lambda idx: idx.split(':')[0]).sum()

# 将分组后的索引替换为每个ID对应的第一个原始索引
grouped.index = df.groupby(lambda idx: idx.split(':')[0]).apply(lambda g: g.index[0])

print(grouped)

最终输出 #

column1  column2
100: Test        3        6
101: FOO         8       16
102: WWW         4        8

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