为多级索引透视表创建按国家分组的条件日期小计

我有一个多级索引透视表，Date为行索引，列层级为[country, Attributes]，生成代码如下：

net_supply = pivot_filter.pivot_table(values='Production', index='Date', columns=['country','Attributes'],aggfunc='sum',fill_value=0)

透视表结构如下：

country1       country2
            prod runs imps prod runs imps
2018-01-01    10   -1    3  5   -6    3
2018-02-01    11   -11   2 13   -2    4
2018-03-01     8    -7    1  1   -5    2
2018-04-01     4    -7    5 10   -7    5
2018-05-01     8   -10    4 11   -2    5

需要为每个国家按Date索引计算条件小计，判断逻辑：

如果 prod + runs >= 0，则小计为 prod + runs + imps；否则仅取 prod + runs

最终仅保留新计算的小计值，理想输出如下：

country1  country2
            subtotal  subtotal 
2018-01-01    12       -1
2018-02-01     2        15
2018-03-01     2        -4
2018-04-01    -3         8
2018-05-01    -2        14

注：部分国家可能仅包含prod、runs等属性中的部分组合，已知margins只能实现普通求和，无法处理条件逻辑，求解决方案。

方法一：按国家列层级分组计算 #

通过对列的country层级分组，针对每个国家的数据单独处理条件逻辑，自动补全缺失属性值：

import pandas as pd
import numpy as np

def calculate_subtotal(country_data):
    # 提取各属性，缺失则补0
    prod = country_data.xs('prod', level='Attributes', axis=1, drop_level=False).fillna(0)
    runs = country_data.xs('runs', level='Attributes', axis=1, drop_level=False).fillna(0)
    imps = country_data.xs('imps', level='Attributes', axis=1, drop_level=False).fillna(0)
    
    # 计算prod + runs的和
    prod_runs = prod + runs
    # 应用条件逻辑生成小计
    subtotal = pd.DataFrame(
        np.where(prod_runs >= 0, prod_runs + imps, prod_runs),
        index=country_data.index,
        columns=[('subtotal',)]
    )
    return subtotal

# 按国家分组计算，再调整列层级顺序
result = net_supply.groupby(level='country', axis=1).apply(calculate_subtotal)
result = result.swaplevel(0, 1, axis=1).sort_index(axis=1)

方法二：重塑数据后透视计算 #

先将多级列转为行结构，统一计算后再透视回目标格式，更直观易读：

import pandas as pd
import numpy as np

# 将列的country层级转为行，方便按国家+日期分组
stacked_data = net_supply.stack(level='country').reset_index()

# 计算prod + runs，缺失属性补0
stacked_data['prod_runs_sum'] = stacked_data['prod'].fillna(0) + stacked_data['runs'].fillna(0)
# 应用条件逻辑计算小计
stacked_data['subtotal'] = np.where(
    stacked_data['prod_runs_sum'] >= 0,
    stacked_data['prod_runs_sum'] + stacked_data['imps'].fillna(0),
    stacked_data['prod_runs_sum']
)

# 重新透视回多级列格式
result = stacked_data.pivot_table(
    values='subtotal',
    index='Date',
    columns=['country'],
    aggfunc='sum'
).rename(columns=lambda c: (c, 'subtotal')).swaplevel(0,1,axis=1).sort_index(axis=1)

两种方法均能处理部分国家属性缺失的场景，最终输出与理想格式完全匹配。

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