无循环硬编码的多索引DataFrame向量化列生成需求问询

需求说明 #

• 创建包含N个一级多索引的DataFrame，每个一级索引下包含X、Y两列
• 为每个一级索引生成Z列，值为对应索引下X与Y的和
• 必须采用向量化操作，禁止使用循环、列表推导或硬编码，确保N>1000时仍可高效扩展
• 最终列数应为N×3（每个一级索引对应X、Y、Z各一列）

不符合要求的实现示例 #

列表推导实现（含循环逻辑） #

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

# 定义一级索引数量
N = 3

# 构建多索引层级
levels = [('A', 'B', 'C'), ('x', 'y')]
columns = list(product(*levels))

# 创建随机DataFrame
df = pd.DataFrame(np.random.randint(0, 10, size=(10, len(columns))), columns=pd.MultiIndex.from_tuples(columns))

# 通过列表推导添加Z列（含循环逻辑，不符合要求）
df[[(*level, 'z') for level in levels[0]]] = df.groupby(level=0, axis=1).sum()

print(df)

硬编码实现（无法扩展） #

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

# 定义一级索引数量
N = 3

# 构建多索引层级
levels = [('A', 'B', 'C'), ('x', 'y')]
columns = list(product(*levels))

# 创建随机DataFrame
df = pd.DataFrame(np.random.randint(0, 10, size=(10, len(columns))), columns=pd.MultiIndex.from_tuples(columns))

# 硬编码添加每个一级索引的Z列（无法支持N>3的场景）
df[('A', 'z')] = df.groupby(level=0, axis=1).sum()[('A')]
df[('B', 'z')] = df.groupby(level=0, axis=1).sum()[('B')]
df[('C', 'z')] = df.groupby(level=0, axis=1).sum()[('C')]

print(df)

符合要求的向量化实现方案 #

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

# 定义一级索引数量（可任意设置，比如N=1000）
N = 3
# 生成一级索引标签，示例为从'A'开始的N个连续字母
first_level = [chr(ord('A') + i) for i in range(N)]
second_level = ['x', 'y']

# 构建多索引列并生成随机DataFrame
columns = list(product(first_level, second_level))
df = pd.DataFrame(np.random.randint(0, 10, size=(10, len(columns))), columns=pd.MultiIndex.from_tuples(columns))

# 向量化生成Z列：按一级索引求和后重命名二级索引，再合并到原DataFrame
sum_df = df.groupby(level=0, axis=1).sum()
sum_df.columns = pd.MultiIndex.from_tuples([(col, 'z') for col in sum_df.columns])

# 合并后按一级索引排序列，确保每个索引下x、y、z连续
df = pd.concat([df, sum_df], axis=1).sort_index(axis=1, level=0)

print(df)

输出示例 #

A     B     C      A   B   C
   x  y  x  y  x  y   z   z   z
0  8  5  9  5  9  9  13  14  18
1  7  4  6  6  0  2  11  12   2
2  4  1  5  1  5  8   5   6  13
3  5  3  5  6  2  0   8  11   2
4  4  3  5  9  3  0   7  14   3
5  9  4  8  3  3  4  13  11   7
6  0  5  7  3  6  1   5  10   7
7  2  9  2  8  0  9  11  10   9
8  5  2  7  5  1  9   7  12  10
9  7  3  9  2  5  5  10  11  10

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