基于列值匹配对两个DataFrame的Sales分组求和（Pandas）

问题背景 #

现有两个Pandas DataFrame：

df1 #

ID  Year Primary_Location Secondary_Location  Sales
0           11  2023          NewYork            Chicago    100
1           11  2023             Lyon      Chicago,Paris    200
2           11  2023           Berlin              Paris    300
3           12  2022          Newyork            Chicago    150
4           12  2022             Lyon      Chicago,Paris    250
5           12  2022           Berlin              Paris    400

df2 #

ID  Year Primary_Location  Sales
0           11  2023          Chicago    150
1           11  2023            Paris    200
2           12  2022          Chicago    300
3           12  2022            Paris    350

需求说明 #

按ID与Year分组，当df2的Primary_Location包含在df1的Secondary_Location中时，将df2对应行的Sales值累加至df1对应行的Sales值。

示例：当ID=11且Year=2023时，Lyon行的Sales需加上df2中Chicago和Paris的Sales值，最终该行Sales为200+150+200=550。

预期输出：

df_primary_output
            ID  Year Primary_Location Secondary_Location  Sales
0           11  2023          NewYork            Chicago    250
1           11  2023             Lyon      Chicago,Paris    550
2           11  2023           Berlin              Paris    500
3           12  2022          Newyork            Chicago    400
4           12  2022             Lyon      Chicago,Paris    900
5           12  2022           Berlin              Paris    750

初始代码 #

import pandas as pd

df1 = pd.DataFrame({'ID': [11, 11, 11, 12, 12, 12],
                   'Year': [2023, 2023, 2023, 2022, 2022, 2022],
                   'Primary_Location': ['NewYork', 'Lyon', 'Berlin', 'Newyork', 'Lyon', 'Berlin'],
                   'Secondary_Location': ['Chicago', 'Chicago,Paris', 'Paris', 'Chicago', 'Chicago,Paris', 'Paris'],
                   'Sales': [100, 200, 300, 150, 250, 400]
                   })

df2 = pd.DataFrame({'ID': [11, 11, 12, 12],
                   'Year': [2023, 2023, 2022, 2022],
                   'Primary_Location': ['Chicago', 'Paris', 'Chicago', 'Paris'],
                   'Sales': [150, 200, 300, 350]
                   })

报错与测试用例 #

运行代码时出现报错：pandas.errors.InvalidIndexError: Reindexing only valid with uniquely valued Index objects，解决方案需同时适配以下测试用例：

测试用例df1 #

Day  ID  Year Primary_Location Secondary_Location  Sales
0       1   11  2023          NewYork            Chicago    100
1       1   11  2023           Berlin            Chicago    300
2       1   11  2022          Newyork            Chicago    150
3       1   11  2022           Berlin            Chicago    400

测试用例df2 #

Day    ID  Year Primary_Location  Sales
0     1     11  2023          Chicago    150
1     1     11  2022          Chicago    300

测试用例预期输出 #

df_primary_output
       Day  ID  Year Primary_Location Secondary_Location  Sales
0       1   11  2023          NewYork            Chicago    250
1       1   11  2023           Berlin            Chicago    450
2       1   11  2022          Newyork            Chicago    450
3       1   11  2022           Berlin            Chicago    700

解决方案 #

核心思路是先按分组键聚合df2的Sales数据，再拆分df1的多地点字段进行匹配累加，最后还原原数据结构，避免索引重复报错。以下是通用可行代码：

import pandas as pd

def calculate_total_sales(df1, df2):
    # 自动识别分组键：取两个DataFrame共有的非位置、非Sales列
    group_keys = [col for col in df1.columns 
                  if col not in ['Primary_Location', 'Secondary_Location', 'Sales'] 
                  and col in df2.columns]
    
    # 按分组键+地点聚合df2的Sales总和
    df2_agg = df2.groupby(group_keys + ['Primary_Location'])['Sales'].sum().reset_index()
    
    # 拆分df1的Secondary_Location为单行单个地点，同时保留原行索引
    df1_expanded = df1.assign(
        Secondary_Location=df1['Secondary_Location'].str.split(',')
    ).explode('Secondary_Location').reset_index(drop=False).rename(columns={'index': 'original_idx'})
    
    # 匹配分组键与地点，合并两个数据集
    merged = pd.merge(
        df1_expanded,
        df2_agg,
        left_on=group_keys + ['Secondary_Location'],
        right_on=group_keys + ['Primary_Location'],
        how='left'
    )
    
    # 按原索引分组，累加匹配到的Sales并保留原数据列
    aggregated = merged.groupby('original_idx').agg(
        **{col: 'first' for col in df1.columns},
        additional_sales=pd.NamedAgg(column='Sales_y', aggfunc=lambda x: x.sum() if x.notna().any() else 0)
    ).reset_index(drop=True)
    
    # 计算最终Sales值并清理临时列
    aggregated['Sales'] = aggregated['Sales'] + aggregated['additional_sales']
    aggregated = aggregated.drop('additional_sales', axis=1)
    
    return aggregated

# 测试原数据集
df_primary_output = calculate_total_sales(df1, df2)
print("原数据集输出：")
print(df_primary_output)

# 测试扩展用例数据集
test_df1 = pd.DataFrame({
    'Day': [1,1,1,1],
    'ID': [11,11,11,11],
    'Year': [2023,2023,2022,2022],
    'Primary_Location': ['NewYork','Berlin','Newyork','Berlin'],
    'Secondary_Location': ['Chicago','Chicago','Chicago','Chicago'],
    'Sales': [100,300,150,400]
})

test_df2 = pd.DataFrame({
    'Day': [1,1],
    'ID': [11,11],
    'Year': [2023,2022],
    'Primary_Location': ['Chicago','Chicago'],
    'Sales': [150,300]
})

test_output = calculate_total_sales(test_df1, test_df2)
print("\n测试用例输出：")
print(test_output)

代码说明 #

1. 自动适配分组键：无需硬编码分组列，可兼容包含Day等额外分组字段的场景。
2. 拆分匹配逻辑：将多地点字段拆分为单行后匹配，确保所有符合条件的地点Sales都被累加。
3. 保留原结构：通过保留原行索引，避免索引重复导致的报错，最终输出与原df1结构完全一致。

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