如何基于列标题扩展Pandas DataFrame的行数？

给定如下分组汇总后的DataFrame：

name  0  1  2
0    A  4  2  1
1    B  2  3  4
2    C  1  3  2

需要将每行数据扩展：除name列外，每行元素根据列标题对应的数值确定保留该列值的行数，超出部分填充NaN，最终得到如下格式的结果：

name 0    1    2
------------------
A    0    1    2
A    0    1    nan
A    0    nan  nan
A    0    nan  nan
B    0    1    2
B    0    1    2
B    nan  1    2
B    nan  nan  2
C    0    1    2
C    nan  1    2
C    nan  1    nan

可复现代码：

import pandas as pd
df = pd.DataFrame({'name': ['A','B','C'], 0:[4,2,1], 1:[2,3,3], 2:[1,4,2]})

之前尝试用df.iterrows()向空DataFrame添加数据，但代码混乱且存在bug，需要更高效简洁的Pandas解决方案。

可以通过构造行掩码矩阵结合repeat和条件赋值实现，避免低效的逐行循环操作：

步骤1：分离列并确定扩展行数 #

先拆分name列与数值列，每行的扩展行数由该行数值的最大值决定：

name_col = df['name']
num_cols = df.drop('name', axis=1)
# 每行需要扩展的行数
row_expand_counts = num_cols.max(axis=1)

步骤2：生成每行的保留掩码 #

对每行数值列，生成对应扩展行数的布尔矩阵：第k行扩展后的记录中，若当前行索引小于该列的数值，则保留列值，否则标记为False：

mask_list = []
for idx, row in num_cols.iterrows():
    expand_count = row_expand_counts[idx]
    # 生成expand_count行的掩码
    row_mask = pd.DataFrame([
        [expand_idx < val for expand_idx, val in enumerate(row)]
        for _ in range(expand_count)
    ])
    mask_list.append(row_mask)

# 合并所有行的掩码并重置索引
full_mask = pd.concat(mask_list, ignore_index=True)

步骤3：扩展name列并填充结果 #

将name列按扩展行数重复，再通过掩码筛选填充数值列（保留位置填列名，否则填pd.NA）：

# 扩展name列
expanded_names = name_col.repeat(row_expand_counts).reset_index(drop=True)

# 构造结果DataFrame
result = pd.DataFrame({
    col: full_mask[col].apply(lambda x: col if x else pd.NA)
    for col in num_cols.columns
})
# 插入name列到首位
result.insert(0, 'name', expanded_names)

完整可运行代码 #

import pandas as pd

# 初始化原始数据
df = pd.DataFrame({'name': ['A','B','C'], 0:[4,2,1], 1:[2,3,3], 2:[1,4,2]})

# 分离列与确定扩展行数
name_col = df['name']
num_cols = df.drop('name', axis=1)
row_expand_counts = num_cols.max(axis=1)

# 生成掩码矩阵
mask_list = []
for idx, row in num_cols.iterrows():
    expand_count = row_expand_counts[idx]
    row_mask = pd.DataFrame([
        [expand_idx < val for expand_idx, val in enumerate(row)]
        for _ in range(expand_count)
    ])
    mask_list.append(row_mask)
full_mask = pd.concat(mask_list, ignore_index=True)

# 构造最终结果
expanded_names = name_col.repeat(row_expand_counts).reset_index(drop=True)
result = pd.DataFrame({
    col: full_mask[col].apply(lambda x: col if x else pd.NA)
    for col in num_cols.columns
})
result.insert(0, 'name', expanded_names)

print(result)

运行输出 #

name     0     1     2
0     A     0     1     2
1     A     0     1  <NA>
2     A     0  <NA>  <NA>
3     A     0  <NA>  <NA>
4     B     0     1     2
5     B     0     1     2
6     B  <NA>     1     2
7     B  <NA>  <NA>     2
8     C     0     1     2
9     C  <NA>     1     2
10    C  <NA>     1  <NA>

（注：Pandas中pd.NA与nan功能等价，显示差异不影响使用）

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