python - 使用多索引在 Pandas 中添加小计列

Question

我有一个数据框，在列上有一个 3 级深度多索引。我想计算跨行 (sum(axis=1)) 的小计，其中我对其中一个级别求和，同时保留其他级别。我想我知道如何使用 pd.DataFrame.sum 的 level 关键字参数来做到这一点。但是，我在考虑如何将这个总和的结果合并回原始表时遇到了麻烦。

设置:

import numpy as np
import pandas as pd
from itertools import product

np.random.seed(0)

colors = ['red', 'green']
shapes = ['square', 'circle']
obsnum = range(5)

rows = list(product(colors, shapes, obsnum))
idx = pd.MultiIndex.from_tuples(rows)
idx.names = ['color', 'shape', 'obsnum']

df = pd.DataFrame({'attr1': np.random.randn(len(rows)), 
                   'attr2': 100 * np.random.randn(len(rows))},
                  index=idx)

df.columns.names = ['attribute']

df = df.unstack(['color', 'shape'])

像这样给出一个漂亮的框架:

假设我想降低 shape 级别。我可以运行:

tots = df.sum(axis=1, level=['attribute', 'color'])

像这样得到我的总数:

有了这个之后，我想将它附加到原始框架上。我想我可以用一种有点麻烦的方式来做到这一点:

tots = df.sum(axis=1, level=['attribute', 'color'])
newcols = pd.MultiIndex.from_tuples(list((i[0], i[1], 'sum(shape)') for i in tots.columns))
tots.columns = newcols
bigframe = pd.concat([df, tots], axis=1).sort_index(axis=1)

有没有更自然的方法来做到这一点？

Answer 1

这是一种没有循环的方法:

s = df.sum(axis=1, level=[0,1]).T
s["shape"] = "sum(shape)"
s.set_index("shape", append=True, inplace=True)
df.combine_first(s.T)

诀窍是使用转置和。因此，我们可以插入另一列(即行)，其中包含附加级别的名称，我们将其命名为与我们求和的名称完全相同。可以使用 set_index 将该列转换为索引中的级别。然后我们将 df 与转置和结合起来。如果总和级别不是最后一个级别，您可能需要对级别进行重新排序。