如何在MultiIndex DataFrame中插入'l'行以补全总计差值？

问题场景 #

现有一个以date和size为多层索引的DataFrame，需要为每个日期插入一行size='l'的数据，该行的volume和hours等于对应日期的total行值减去s和m行值的和。原始数据构造代码如下：

import numpy as np
import pandas as pd

date = pd.date_range('2023-01-01', periods=3)
size = ['s','m','total']
arrays = [date, size]
index = pd.MultiIndex.from_product(arrays, names=['date','size'])
volume = [4,15,47,8,12,46,4,14,48]
hours = [2,1,13,4,4,10,1,2,10]
df = pd.DataFrame({'volume':volume, 'hours':hours}, index=index)

解决方案 #

可以通过分组计算、合并数据再重新排序索引来实现，具体代码如下：

# 定义函数，计算单个日期分组下的l行数据
def get_l_row(group):
    total_vals = group.xs('total')
    s_m_sum = group.loc[['s', 'm']].sum()
    l_vals = total_vals - s_m_sum
    l_vals.name = 'l'
    return pd.DataFrame(l_vals).T

# 按date分组计算所有l行数据，并构建匹配的MultiIndex
l_rows = df.groupby(level='date').apply(get_l_row)
l_rows.index = pd.MultiIndex.from_tuples(
    [(date, 'l') for date in l_rows.index.get_level_values(0)],
    names=['date', 'size']
)

# 合并原数据与l行数据，再调整每个date下的size顺序
combined_df = pd.concat([df, l_rows])
desired_order = ['s', 'm', 'l', 'total']
final_df = combined_df.groupby(level='date').apply(
    lambda g: g.reindex(desired_order, level='size')
).droplevel(0)

print(final_df)

输出结果 #

运行代码后得到符合要求的DataFrame：

volume  hours
date       size
2023-01-01 s           4      2
           m          15      1
           l          28     10
           total      47     13
2023-01-02 s           8      4
           m          12      4
           l          26      2
           total      46     10
2023-01-03 s           4      1
           m          14      2
           l          30      7
           total      48     10

代码说明 #

1. get_l_row函数：针对单个日期分组，提取total行数据，计算s和m行的求和值，最终得到l行的计算结果。
2. 构建l_rows的索引：保证新行的多层索引格式与原数据一致，避免合并时出现索引不匹配问题。
3. 调整顺序：通过groupby+reindex固定每个日期下size的展示顺序，确保输出符合预期的s→m→l→total排列。

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