如何在Pandas中针对主键列重复值将数据填充至相邻列

问题描述 #

初始表格：

+---------+------+------+------+------+------+------+------+------+
| COL_ID  | ColB | COLC | COLD | COLE | COLF | COLG | COLH | COLI |
+---------+------+------+------+------+------+------+------+------+
| aa1     |    1 |    1 |      |      |      |      |      |      |
| aa1     | 2    |    1 |      |      |      |      |      |      |
| aa2     |    3 |    1 |      |      |      |      |      |      |
| ab3     |    6 |    2 |      |      |      |      |      |      |
| ab3     |    5 |    2 |      |      |      |      |      |      |
| ab3     |    7 |    1 |      |      |      |      |      |      |
| ab3     |    1 |    1 |      |      |      |      |      |      |
+---------+------+------+------+------+------+------+------+------+

需求：当主键列COL_ID出现重复值时，将分组内每行的ColB、COLC值按顺序依次填充至后续的空白列对（COLD&COLE、COLF&COLG、COLH&COLI），最终每个COL_ID仅保留一行，得到如下结果：

+---------+------+------+------+------+------+------+------+------+
| COL_ID  | ColB | COLC | COLD | COLE | COLF | COLG | COLH | COLI |
+---------+------+------+------+------+------+------+------+------+
| aa1     |    1 |    1 |    2 |    1 |      |      |      |      |
| aa2     |    3 |    1 |      |      |      |      |      |      |
| ab3     |    6 |    2 |    5 |    2 |    7 |    1 |    1 |    1 |
+---------+------+------+------+------+------+------+------+------+

注：原示例目标表格中aa1行的COLD、COLE值存在笔误，上述结果为按逻辑推导的正确输出。

初始数据生成代码 #

import pandas as pd
import numpy as np

# 修正原代码笔误：补充COLE列，修正COLI的键名
my_dic =  {'COL_ID': ['aa1', 'aa1', 'aa2', 'ab3','ab3','ab3','ab3'],
           'COLB': [1,2,3,6,5,7,1],
           'COLC': [1,1,1,2,2,1,1],
           'COLD':[np.nan]*7,
           'COLE':[np.nan]*7,
           'COLF':[np.nan]*7,
           'COLG':[np.nan]*7,
           'COLH':[np.nan]*7,
           'COLI':[np.nan]*7}
    
df = pd.DataFrame(my_dic)

解决方案代码 #

# 定义分组处理函数：将每组的COLB、COLC按行顺序扁平化
def process_group(group):
    # 提取每组的COLB和COLC值，按行拼接成一维数组
    flat_values = group[['COLB', 'COLC']].values.flatten()
    # 定义所有目标列的顺序
    all_columns = ['COLB', 'COLC', 'COLD', 'COLE', 'COLF', 'COLG', 'COLH', 'COLI']
    # 截取与扁平化数组长度匹配的列名
    target_columns = all_columns[:len(flat_values)]
    # 返回Series，自动对齐列名
    return pd.Series(flat_values, index=target_columns)

# 按COL_ID分组处理，重置索引并对齐原表格列
result = df.groupby('COL_ID').apply(process_group).reset_index()
result = result.reindex(columns=df.columns)

# 打印结果
print(result.to_string(index=False))

代码说明 #

1. 修正初始数据：补充原代码遗漏的COLE列，修正COLI:的键名错误，确保列名与目标表格一致。
2. 分组处理逻辑：
   • 按COL_ID分组后，将每组内的COLB和COLC按行顺序拼接成一维数组，例如ab3分组会生成[6,2,5,2,7,1,1,1]。
   • 按照目标列的顺序，将扁平化数组对应填充到列中，多余的列保持为空（NaN）。
3. 结果对齐：通过reindex确保结果列顺序与原表格完全一致。

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