Pandas pivot_table 高级应用：生成多级列标题及全组合小计

import pandas as pd
from itertools import combinations

data = {
    'date': ['Jan', 'Feb'],
    'subject': ['English', 'Chemistry'],
    'animal': ['cat', 'dog'],
    'colors': ['blue', 'green'],
    'value': [1, 2]
}
df = pd.DataFrame(data)

# 扩展数据以展示更多组合，确保 'all' 类别有数据可聚合
df_expanded = pd.DataFrame([
    ['Jan', 'English', 'cat', 'blue', 1],
    ['Feb', 'Chemistry', 'dog', 'green', 2],
    ['Jan', 'English', 'dog', 'blue', 3], # 新增数据
    ['Feb', 'Chemistry', 'cat', 'green', 4], # 新增数据
    ['Jan', 'English', 'cat', 'green', 5], # 新增数据
], columns=['date', 'subject', 'animal', 'colors', 'value'])

print("原始数据框:")
print(df_expanded)

cols_to_categorize = ['date', 'subject', 'animal', 'colors']
cats = {col: pd.CategoricalDtype(list(df_expanded[col].unique()) + ['all'], ordered=True)
        for col in cols_to_categorize}
df_categorized = df_expanded.astype(cats)

print("\n转换为CategoricalDtype后的数据框:")
print(df_categorized.dtypes)

# 存储所有数据（原始数据 + 各种小计数据）
all_data_for_pivot = [df_categorized]

# 遍历所有 N-1 列的组合，生成各种小计行
# 例如，对于 ['date', 'subject', 'animal', 'colors']，会生成以下组合：
# (date, subject, animal) -> colors 列为 'all'
# (date, subject, colors) -> animal 列为 'all'
# ...
for r in range(1, len(cols_to_categorize) + 1):
    for grp_cols in combinations(cols_to_categorize, r=r):
        if len(grp_cols) == len(cols_to_categorize): # 跳过全列组合，因为这已是原始数据
            continue

        # 对当前组合的列进行分组求和
        df_subtotal = df_categorized.groupby(list(grp_cols), as_index=False, observed=True)['value'].sum()

        # 对于未参与分组的列，将其值设为 'all'
        missing_cols = [col for col in cols_to_categorize if col not in grp_cols]
        for m_col in missing_cols:
            df_subtotal[m_col] = 'all'
            # 确保 'all' 类别被正确识别
            df_subtotal[m_col] = df_subtotal[m_col].astype(cats[m_col])

        all_data_for_pivot.append(df_subtotal)

# 合并所有数据，包括原始数据和各种小计数据
# 使用 pd.concat 合并时，如果某个分类列在某个DataFrame中缺失，会填充为 NaN，
# 但由于我们已经用 'all' 填充了缺失的类别，这里主要是合并不同的聚合结果。
# 最终的 DataFrame 将包含所有原始数据点以及各种小计数据点，其中某些分类列的值为 'all'。
df_final_prep = pd.concat(all_data_for_pivot, ignore_index=True)

print("\n预处理后的数据框（部分示例行，包含'all'类别）:")
print(df_final_prep.head(10)) # 打印前10行查看效果

# 定义透视表的行、列和值
index_col = 'date'
columns_cols = ['subject', 'animal', 'colors']
values_col = 'value'
agg_func = 'sum' # 聚合函数，可以根据需求修改为 'median' 等

# 执行 pivot_table 操作
# observed=False 是关键，它会强制 pivot_table 考虑 CategoricalDtype 中定义的所有类别，
# 即使某些组合在 df_final_prep 中没有直接对应的行。
# fill_value=-1 用于填充那些在原始数据中不存在但由 'all' 类别组合生成的新单元格。
result_df = df_final_prep.pivot_table(
    index=index_col,
    columns=columns_cols,
    values=values_col,
    aggfunc=agg_func,
    fill_value=0, # 填充值为0，更符合小计的语义
    observed=False
)

print("\n最终的多级列标题小计数据框:")
print(result_df)