Pandas DataFrame基于共有员工号合并关联记录的递归方案求解

可行性判断 #

递归遍历Pandas DataFrame的方式理论可实现但完全不适合落地：该方式需要逐行逐员工号对比公共ID，时间复杂度最高可达O(n²)，数据量超过百级别性能就会严重下滑，同时递归深度过高会触发Python递归深度限制报错，实用性极差。

落地解决方案 #

该需求本质是连通分量计算问题：只要两行存在公共员工号，就属于同一分组，最终将分组内所有员工号合并后赋值给组内所有行，使用并查集（Union-Find） 数据结构可以高效实现，具体代码如下：

步骤1：数据预处理 #

import pandas as pd
from collections import defaultdict

# 构造示例输入
df = pd.DataFrame({
    "Company": [1, 2, 3, 4],
    "Employee Number": ["12", "34, 12", "56, 34, 78", "90"]
})

# 将员工号拆分转为集合，方便后续计算
df["emp_set"] = df["Employee Number"].str.split(",\s*").apply(set)

步骤2：用并查集计算分组 #

# 简易并查集实现
class UnionFind:
    def __init__(self, size):
        self.parent = list(range(size))
    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]
    def union(self, x, y):
        x_root, y_root = self.find(x), self.find(y)
        if x_root != y_root:
            self.parent[y_root] = x_root

uf = UnionFind(len(df))
emp_index_map = defaultdict(list)
# 记录每个员工号对应的所有行索引
for idx, emp_set in df["emp_set"].items():
    for emp in emp_set:
        emp_index_map[emp].append(idx)
# 将同员工号对应的所有行合并到同一分组
for indices in emp_index_map.values():
    if len(indices) < 2:
        continue
    base_idx = indices[0]
    for idx in indices[1:]:
        uf.union(base_idx, idx)
# 给每行标记分组ID
df["group_id"] = [uf.find(i) for i in range(len(df))]

步骤3：聚合得到最终结果 #

# 按分组聚合所有不重复员工号，排序后转为要求的字符串格式
group_emp = df.groupby("group_id")["emp_set"].agg(lambda x: set.union(*x)).reset_index(name="all_emp")
group_emp["all_emp"] = group_emp["all_emp"].apply(lambda x: ", ".join(sorted(x, key=int)))
# 合并回原表得到输出
result = df.merge(group_emp, on="group_id")[["Company", "all_emp"]].rename(columns={"all_emp": "Employee Number"})

输出结果 #

最终result变量值与预期完全一致：

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