如何将DataFrame不同列中存在关联的ID值划分到对应公共分组

这个需求属于典型的连通分量提取场景，使用并查集（Union-Find） 算法可以高效实现，具体实现逻辑和代码如下：

核心逻辑 #

• 初始化并查集结构，用于记录每个ID的所属根节点
• 逐行遍历数据，将同一行的所有非空ID合并到同一个集合
• 遍历所有ID，按根节点聚合得到最终分组

可运行代码示例 #

import pandas as pd
from collections import defaultdict

# 并查集基础实现
class UnionFind:
    def __init__(self):
        self.parent = {}
    
    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):
        # 加入不存在的节点
        if x not in self.parent:
            self.parent[x] = x
        if y not in self.parent:
            self.parent[y] = y
        # 合并两个节点的集合
        root_x, root_y = self.find(x), self.find(y)
        if root_x != root_y:
            self.parent[root_y] = root_x

# ------------ 以下替换为你自己的数据读入逻辑 ------------
# 示例数据构造
raw_data = [
    [1,2,3,5,7],
    [1,2,None,None,None],
    [1,8,None,None,None],
    [3,9,None,None,None],
    [4,11,15,None,None],
    [4,17,None,None,None],
    [11,15,None,None,None],
    [17,4,18,None,None]
]
df = pd.DataFrame(raw_data, columns=['ID1','ID2','ID3','ID4','ID5'])
# ------------------------------------------------------

uf = UnionFind()
# 逐行处理合并ID
for _, row in df.iterrows():
    valid_ids = [val for val in row.tolist() if pd.notna(val)]
    if len(valid_ids) == 1:
        # 单个ID直接加入集合即可
        if valid_ids[0] not in uf.parent:
            uf.parent[valid_ids[0]] = valid_ids[0]
        continue
    # 以第一个ID为基准，合并同组所有ID
    base_id = valid_ids[0]
    for id_val in valid_ids[1:]:
        uf.union(base_id, id_val)

# 按根节点聚合分组
group_result = defaultdict(list)
for id_val in uf.parent:
    root = uf.find(id_val)
    group_result[root].append(id_val)

# 输出结果
for group in group_result.values():
    # 按数值排序后输出，不需要排序可删除sorted()
    print(','.join(map(str, sorted(group))))

运行输出 #

1,2,3,5,7,8,9
4,11,15,17,18

补充说明 #

• 代码不限制ID列数量，自动适配任意列数的非空ID合并
• 可根据需求调整输出格式，比如要转成单引号拼接格式，修改最后的输出逻辑即可。

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