pandas dataframe按idrun分组对指定idbasin的q值求和合并的实现方法

以下方案全程使用pandas内置算子，无自定义循环，针对60万行规模的数据可以在毫秒级完成处理：

1. 先拆分数据集为两部分：不需要参与合并的行、需要合并计算的行
2. 对需要合并的行按idrun、time分组聚合，计算q的总和，同时将idbasin统一设置为给定元组的第一个元素
3. 拼接两部分数据后按业务逻辑排序即可

示例代码 #

import pandas as pd

# 输入参数：待合并的idbasin元组
basin_tuple = (1, 2)
target_basin = basin_tuple[0]

# 拆分数据
df_to_merge = df[df['idbasin'].isin(basin_tuple)].copy()
df_keep = df[~df['idbasin'].isin(basin_tuple)].copy()

# 聚合计算合并后的行
df_merged = df_to_merge.groupby(['idrun', 'time'], as_index=False).agg(
    idbasin = ('idbasin', lambda x: target_basin),
    q = ('q', 'sum'),
    # 剩余的4列可根据你的业务需求补充聚合规则，示例如下：
    # 列名1 = ('原列名1', 'first'), # 取分组内第一个值
    # 列名2 = ('原列名2', 'sum'), # 数值列求和
)

# 合并结果并排序
df_result = pd.concat([df_keep, df_merged], ignore_index=True)
df_result = df_result.sort_values(by=['idrun', 'idbasin', 'time'], ignore_index=True)

补充说明 #

如果你的其余列的取值完全由idrun+idbasin+time唯一确定，且合并后的新行要沿用原idbasin=target_basin对应的其他列取值，可以先把待合并部分里idbasin等于目标值的其他列取出来，再和求和后的q关联即可，无需修改聚合规则。

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