在Pandas中创建包含所有组合小计的多层级透视表

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)

# 扩展数据以展示更复杂的组合，虽然原始问题只提供了2行
# 这里为了演示效果，我们可以稍微扩展一下数据
extended_data = {
    'date': ['Jan', 'Feb', 'Jan', 'Feb', 'Mar'],
    'subject': ['English', 'Chemistry', 'English', 'Math', 'Chemistry'],
    'animal': ['cat', 'dog', 'dog', 'cat', 'cat'],
    'colors': ['blue', 'green', 'blue', 'red', 'green'],
    'value': [1, 2, 3, 4, 5]
}
df = pd.DataFrame(extended_data)

print("原始DataFrame:")
print(df)

cols_to_categorize = ['date', 'subject', 'animal', 'colors']
cats = {}
for col in cols_to_categorize:
    # 获取列的唯一值，并添加 'all' 作为新的类别
    unique_values = df[col].unique().tolist()
    cats[col] = pd.CategoricalDtype(unique_values + ['all'], ordered=True)

# 将DataFrame的相应列转换为这些分类类型
df_categorized = df.astype(cats)

print("\n转换为CategoricalDtype后的DataFrame (部分):")
print(df_categorized.dtypes)

intermediate_data_frames = [df_categorized] # 包含原始数据

# 生成所有可能的分组组合 (例如，对N-1个列进行分组)
# 这里的策略是，通过聚合部分列，然后与原始数据合并，
# 缺失的列会变成NaN，再用'all'填充，从而创建出'all'的组合行
for r in range(1, len(cols_to_categorize)):
    for grp_cols in combinations(cols_to_categorize, r=r):
        # 聚合指定列，并对'value'求和（或median等）
        grouped_df = df_categorized.groupby(list(grp_cols), as_index=False, observed=True)['value'].sum()
        intermediate_data_frames.append(grouped_df)

# 合并所有数据帧，并用 'all' 填充因聚合而产生的 NaN 值
# 这里的关键是当grouped_df与df_categorized合并时，grouped_df中缺少的分类列会产生NaN，
# 这些NaN随后被'all'填充，从而构建出'all'的组合
combined_df_with_subtotals = pd.concat(intermediate_data_frames, ignore_index=True)
combined_df_with_subtotals = combined_df_with_subtotals.fillna('all')

# 确保填充后的列仍然是正确的CategoricalDtype
for col in cols_to_categorize:
    combined_df_with_subtotals[col] = combined_df_with_subtotals[col].astype(cats[col])

print("\n包含'all'小计的中间数据集 (部分):")
print(combined_df_with_subtotals.head())
print(combined_df_with_subtotals.tail())

# 执行pivot_table操作
# index='date' 作为行索引
# columns=['subject', 'animal', 'colors'] 作为多层级列索引
# values='value' 是要聚合的值
# aggfunc='sum' 是聚合函数 (可以根据需求改为 'median', 'mean' 等)
# fill_value=-1 填充数据中不存在的组合值 (例如，某个日期没有某个组合的数据)
# observed=False 确保所有定义的分类组合都会作为列显示
final_pivot_table = combined_df_with_subtotals.pivot_table(
    index='date',
    columns=['subject', 'animal', 'colors'],
    values='value',
    aggfunc='sum', # 原始问题是median，但答案用sum，这里保持sum，可根据需求修改
    fill_value=0, # 填充缺失值，通常用0或NaN
    observed=False # 确保所有分类组合都出现
)

print("\n最终的多层级透视表，包含所有组合小计:")
print(final_pivot_table)