如何为Pandas数据框中POS长度超2的行标记对应根节点KOMPID

示例代码与初始数据 #

导入pandas库并创建DataFrame：

import pandas as pd
collection   = ({
    'REZID':["E0912","E0912","E0912","E0912","E0912","E0912","E0912","E0912","E0913","E0913",],
    'POS' :["01","0101","0102","0103","010301","010302","02","0201","01","0101"],
    'KOMPID':['k01','k02','k03','k04','k05','k06','k07','k08','k09','k10'],
    'WEIGHT':[1000,300,400,300,150,150,1400,1400,1200,500]
                })
df = pd.DataFrame(collection, columns=['REZID','POS','KOMPID','WEIGHT'])

初始DataFrame输出：

REZID     POS KOMPID  WEIGHT
0  E0912      01    k01    1000
1  E0912    0101    k02     300
2  E0912    0102    k03     400
3  E0912    0103    k04     300
4  E0912  010301    k05     150
5  E0912  010302    k06     150
6  E0912      02    k07    1400
7  E0912    0201    k08    1400
8  E0913      01    k09    1200
9  E0913    0101    k10     500

需求说明 #

需要为POS字段长度超过2位的行标记对应根节点，根节点为相同REZID下POS长度为2位的行的KOMPID值。已通过以下代码筛选出根节点数据：

root = df.loc[df['POS'].str.len() == 2]

根节点数据输出：

REZID POS KOMPID  WEIGHT
0  E0912  01    k01    1000
6  E0912  02    k07    1400
8  E0913  01    k09    1200

解决方案 #

方法1：使用映射字典（简洁直观） #

先构建根节点的映射关系，再通过apply为每行匹配根节点：

# 构建(REZID, 根POS)到KOMPID的映射字典
root_map = root.set_index(['REZID', 'POS'])['KOMPID'].to_dict()
# 提取每行POS的前2位，结合REZID匹配根节点
df['ROOT'] = df.apply(lambda row: root_map[(row['REZID'], row['POS'][:2])], axis=1)

方法2：使用合并操作（高效适合大数据量） #

避免apply的循环操作，通过字段提取+表合并实现：

# 提取POS字段前2位作为根节点的POS标识
df['ROOT_POS'] = df['POS'].str[:2]
# 与根节点数据合并，匹配对应KOMPID
df = df.merge(root[['REZID', 'POS', 'KOMPID']], 
              left_on=['REZID', 'ROOT_POS'], 
              right_on=['REZID', 'POS'], 
              how='left')
# 重命名列并清理临时字段
df = df.rename(columns={'KOMPID_y': 'ROOT'}).drop(columns=['POS_y', 'ROOT_POS'])
# 为根节点自身填充ROOT值为自身KOMPID
df['ROOT'] = df['ROOT'].fillna(df['KOMPID'])

最终结果 #

执行上述任意一种方法后，得到添加ROOT列的DataFrame：

REZID     POS KOMPID  WEIGHT ROOT
0  E0912      01    k01    1000  k01
1  E0912    0101    k02     300  k01
2  E0912    0102    k03     400  k01
3  E0912    0103    k04     300  k01
4  E0912  010301    k05     150  k01
5  E0912  010302    k06     150  k01
6  E0912      02    k07    1400  k07
7  E0912    0201    k08    1400  k07
8  E0913      01    k09    1200  k09
9  E0913    0101    k10     500  k09

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