如何基于现有数据在多层索引DataFrame中新增聚合行

全程用pandas向量化操作实现，无循环，大数据量下运行效率很高，扩展核心国家只需要修改对应列表即可。

步骤1：导入依赖、加载数据 #

import pandas as pd

# 以下为示例数据构造，实际使用时替换为你自己的多层索引DataFrame即可
data = [
    ['USA', 'DEU', 20],
    ['USA', 'CHN', 10],
    ['USA', 'MEX', 5],
    ['DEU', 'USA', 12],
    ['DEU', 'CHN', 6],
    ['DEU', 'MEX', 2],
    ['CHN', 'USA', 1],
    ['CHN', 'DEU', 2],
    ['CHN', 'MEX', 3],
]
df = pd.DataFrame(data, columns=['Source', 'Partner', 'Value']).set_index(['Source', 'Partner'])

# 定义核心国家列表，后续调整核心国家范围只需要修改这个列表
core_countries = ['USA', 'DEU', 'CHN']

步骤2：提取核心国家间流量 #

聚合统计仅计算核心国家之间的流量，不需要引入非核心国家的双边数据：

core_mask = (
    df.index.get_level_values('Source').isin(core_countries)
    & df.index.get_level_values('Partner').isin(core_countries)
)
core_flow = df[core_mask]

步骤3：计算两类聚合行 #

# 计算Source为CORE的聚合行：统计所有核心国作为来源，流向每个核心国的总价值
agg_source_core = core_flow.groupby('Partner')['Value'].sum().to_frame()
agg_source_core.index = pd.MultiIndex.from_product(
    [['CORE'], agg_source_core.index],
    names=['Source', 'Partner']
)

# 计算Partner为CORE的聚合行：统计每个核心国作为来源，流向所有核心国的总价值
agg_partner_core = core_flow.groupby('Source')['Value'].sum().to_frame()
agg_partner_core.index = pd.MultiIndex.from_product(
    [agg_partner_core.index, ['CORE']],
    names=['Source', 'Partner']
)

步骤4：拼接得到最终结果 #

# 保留原表所有行，新增聚合行后按索引排序
result = pd.concat([df, agg_source_core, agg_partner_core]).sort_index()

结果说明 #

运行后得到的结果中，新增聚合行的数值和预期完全匹配：

Value
Source       Partner              
CHN          CORE          3
             DEU           2
             MEX           3
             USA           1
CORE         CHN          16
             DEU          22
             USA          13
DEU          CHN           6
             CORE         18
             MEX           2
             USA          12
USA          CHN          10
             CORE         30
             DEU          20
             MEX           5

如果不需要保留原表中非核心国家的相关行，拼接时把原df替换为core_flow即可。

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