如何在Pandas多级列DataFrame中设置mass数组为索引

初始代码与输出 #

以下是生成多级列DataFrame的代码：

import pandas as pd
import numpy as np
from collections import defaultdict

np.random.seed(0)

pd_df = defaultdict(list)
categories = ['cat 1', 'cat 2']
sig_bkgd = ['signal', 'bkgd']
masses = [i for i in range(10, 20)]

for m in masses:
    for cat in categories:
        for sb in sig_bkgd:
            pd_df[(cat, sb)].append(np.random.randint(100))

pd.DataFrame(pd_df)

对应的输出（保留多级列结构）：

cat 1       cat 2     
  signal bkgd signal bkgd
0     44   47     64   67
1     67    9     83   21
2     36   87     70   88
3     88   12     58   65
4     39   87     46   88
5     81   37     25   77
6     72    9     20   80
7     69   79     47   64
8     82   99     88   49
9     29   19     19   14

需求 #

希望将masses数组设置为DataFrame的索引，同时保留上述的多级列结构。

尝试的方法及问题 #

方法1：添加普通列后设为索引 #

代码：

pd_df['Mass'] = masses
for m in masses:
    for cat in categories:
        for sb in sig_bkgd:
            pd_df[(cat, sb)].append(np.random.randint(100))

pd.DataFrame(pd_df).set_index("Mass")

问题：生成的DataFrame丢失了多级列结构，列变成了元组形式：

(cat 1, signal)  (cat 1, bkgd)  (cat 2, signal)  (cat 2, bkgd)
Mass                                                                 
10                 44             47               64             67
11                 67              9               83             21
12                 36             87               70             88
13                 88             12               58             65
14                 39             87               46             88
15                 81             37               25             77
16                 72              9               20             80
17                 69             79               47             64
18                 82             99               88             49
19                 29             19               19             14

方法2：添加重复的多级列作为索引 #

代码：

for m in masses:
    for cat in categories:
        pd_df[(cat, 'mass')].append(m)
        for sb in sig_bkgd:
            pd_df[(cat, sb)].append(np.random.randint(100))

pd.DataFrame(pd_df).set_index([('cat 1', 'mass'), ('cat 2', 'mass')])

问题：索引出现冗余，包含两个重复的mass层级：

cat 1       cat 2     
                            signal bkgd signal bkgd
(cat 1, mass) (cat 2, mass)                        
10            10                44   47     64   67
11            11                67    9     83   21
12            12                36   87     70   88
13            13                88   12     58   65
14            14                39   87     46   88
15            15                81   37     25   77
16            16                72    9     20   80
17            17                69   79     47   64
18            18                82   99     88   49
19            19                29   19     19   14

解决方案 #

有两种简单的方式可以实现需求：

方式1：先创建DataFrame，再设置索引 #

先按原逻辑生成多级列的DataFrame，之后直接将masses赋值给索引，无需修改字典：

import pandas as pd
import numpy as np
from collections import defaultdict

np.random.seed(0)

pd_df = defaultdict(list)
categories = ['cat 1', 'cat 2']
sig_bkgd = ['signal', 'bkgd']
masses = [i for i in range(10, 20)]

# 原逻辑生成数据
for m in masses:
    for cat in categories:
        for sb in sig_bkgd:
            pd_df[(cat, sb)].append(np.random.randint(100))

# 创建DataFrame后设置索引
df = pd.DataFrame(pd_df)
df.index = masses
# 可选：给索引命名
df.index.name = 'Mass'

print(df)

输出结果（保留多级列结构，索引为Mass）：

cat 1       cat 2     
    signal bkgd signal bkgd
Mass                        
10       44   47     64   67
11       67    9     83   21
12       36   87     70   88
13       88   12     58   65
14       39   87     46   88
15       81   37     25   77
16       72    9     20   80
17       69   79     47   64
18       82   99     88   49
19       29   19     19   14

方式2：创建DataFrame时指定索引 #

在pd.DataFrame()中直接通过index参数设置：

# 生成数据的逻辑同上...
df = pd.DataFrame(pd_df, index=masses)
df.index.name = 'Mass'

效果和方式1完全一致。

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