高效生成Pandas DataFrame附加索引的方法

我正在处理一个大型的二进制（n-hot编码）DataFrame，示例结构如下：

import pandas as pd

data = {
  'A' : [0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
  'B' : [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
  'C' : [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1],
  'D' : [1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0],
  'E' : [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
}

df = pd.DataFrame.from_dict(data)

原始DataFrame输出：

A  B  C  D  E
0   0  0  0  1  0
1   0  0  0  1  0
2   0  0  0  1  0
3   1  0  0  0  0
4   1  0  0  0  0
5   1  0  0  0  0
6   1  0  0  0  0
7   1  0  0  0  0
8   0  0  0  1  0
9   0  0  0  1  0
10  0  0  1  0  0
11  0  0  1  0  0
12  0  0  1  0  0
13  0  0  1  0  0

为提升数据提取与检索效率，需要生成三类附加索引：

• 标识每行中包含1的列名
• 该组的批次序号（连续相同列名的行视为一个批次）
• 批次内的行序号

期望最终输出：

A  B  C  D  E
0   1  1  D  0  0  0  1  0
1   1  2  D  0  0  0  1  0
2   1  3  D  0  0  0  1  0
3   1  1  A  1  0  0  0  0
4   1  2  A  1  0  0  0  0
5   1  3  A  1  0  0  0  0
6   1  4  A  1  0  0  0  0
7   1  5  A  1  0  0  0  0
8   2  1  D  0  0  0  1  0
9   2  2  D  0  0  0  1  0
10  1  1  C  0  0  1  0  0
11  1  2  C  0  0  1  0  0
12  1  3  C  0  0  1  0  0
13  1  4  C  0  0  1  0  0

高效实现方案 #

针对大型DataFrame，必须使用矢量化操作避免循环，以下是具体步骤：

1. 提取包含1的列名 #

示例中每行仅存在一个1，直接用idxmax高效获取；若为多1场景，可改为df.apply(lambda x: ','.join(df.columns[x==1]), axis=1)：

# 生成标识列（col_name）
df['col_name'] = df.idxmax(axis=1)

2. 生成批次序号 #

通过对比当前行与上一行的col_name是否变化，生成分组标记后累加得到批次号：

# 生成批次序号（batch_num）
df['batch_num'] = (df['col_name'] != df['col_name'].shift(1)).cumsum()

3. 生成批次内序号 #

按col_name和batch_num分组后，用cumcount()生成组内行号（+1转为从1开始）：

# 生成批次内序号（inner_num）
df['inner_num'] = df.groupby(['col_name', 'batch_num']).cumcount() + 1

4. 调整列顺序与设置索引（可选） #

若需要将附加列前置或设为复合索引，可执行：

# 调整列顺序，将附加列移至前面
df = df[['batch_num', 'inner_num', 'col_name'] + list(df.columns[:-3])]
# 设置复合索引（按需选择，优化检索效率）
df = df.set_index(['batch_num', 'inner_num', 'col_name'])

执行上述代码后，即可得到符合需求的结果。

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