如何高效将Pandas计数行转换为独热编码？

问题背景 #

原始DataFrame格式：

A  B  C  D  E
0  a  2  0  3  x
1  b  1  2  0  y

需要转换为每个计数值对应一行独热编码的格式：

A  B  C  D  E
0  a  1  0  0  x
1  a  1  0  0  x
2  a  0  0  1  x
3  a  0  0  1  x
4  a  0  0  1  x
5  b  1  0  0  y
6  b  0  1  0  y
7  b  0  1  0  y

用户原实现代码（低效）：

# Sample DataFrame
data = {
    'A': ['a', 'b'],
    'B': [2, 1],
    'C': [0, 2],
    'D': [3, 0],
    'E': ['x', 'y']
}
df = pd.DataFrame(data)

new_df = pd.DataFrame(columns=df.columns)

for index, row in df.iterrows():
    first_val = row.iloc[0]
    last_val = row.iloc[-1]
    middle_vals = row.iloc[1:-1].astype(int)
    for i in range(len(middle_vals)):
        new_data = [first_val] + [1 if i == j else 0 for j in range(len(middle_vals))] + [last_val]
        new_rows = pd.DataFrame([new_data] * middle_vals.iloc[i], columns=df.columns)
        new_df = pd.concat([new_df, new_rows], ignore_index=True)

问题：逐行循环+多次concat导致效率极低，寻求向量化优化方案。

向量化优化实现 #

核心是利用Pandas内置的向量化操作替代循环，避免多次拼接开销：

import pandas as pd

data = {
    'A': ['a', 'b'],
    'B': [2, 1],
    'C': [0, 2],
    'D': [3, 0],
    'E': ['x', 'y']
}
df = pd.DataFrame(data)

# 定义首尾列和中间计数列
first_col, last_col = df.columns[0], df.columns[-1]
middle_cols = df.columns[1:-1]

# 1. 将宽表转为长格式，保留首尾列，过滤计数为0的行
melted = df.melt(
    id_vars=[first_col, last_col],
    value_vars=middle_cols,
    var_name='target_col',
    value_name='count'
)
melted = melted[melted['count'] > 0]

# 2. 按计数值重复行（向量化操作）
repeated_rows = melted.loc[melted.index.repeat(melted['count'])]

# 3. 生成独热编码
one_hot = pd.get_dummies(repeated_rows['target_col'], prefix='', prefix_sep='')

# 4. 合并列并恢复原始顺序
result = pd.concat(
    [repeated_rows[[first_col, last_col]].reset_index(drop=True), one_hot],
    axis=1
)[df.columns]

print(result)

优化点说明 #

1. melt转长格式：把多列计数转为单行列名+计数的结构，一次性处理所有行，避免逐行遍历。
2. repeat批量重复：直接对索引进行重复操作，比循环创建小DataFrame效率高几个数量级。
3. get_dummies自动生成独热编码：无需手动判断0/1，内置向量化处理速度快。
4. 单次合并：仅在最后合并一次数据，避免多次concat的内存开销。

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