如何识别Pandas两列重叠关联元素的分组并新增分组标识列

实现方案 #

该需求本质是求解无向图的连通分量：每一行的两个ID对应一条无向边，所有连通的节点属于同一个分组。我们可以通过**并查集（DSU）**算法快速实现，无需额外依赖。

完整代码实现 #

import pandas as pd

# 示例数据
so_fake_data = pd.DataFrame(
    {
        'id_1': ['A', 'A', 'D', 'F', 'H', 'H', 'K', 'M'],
        'id_2': ['B', 'C', 'E', 'G', 'I', 'J', 'L', 'L']
    }
)

# 并查集实现
class DSU:
    def __init__(self, nodes):
        self.parent = {node: node for node in nodes}
    
    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 = self.find(x)
        y_root = self.find(y)
        if x_root != y_root:
            self.parent[y_root] = x_root

# 收集所有出现过的唯一节点
all_nodes = pd.concat([so_fake_data['id_1'], so_fake_data['id_2']]).unique()
# 初始化并查集
dsu = DSU(all_nodes)
# 遍历所有边执行合并操作
for _, row in so_fake_data.iterrows():
    dsu.union(row['id_1'], row['id_2'])

# 生成每个节点对应的分组标识（此处用连通分量的根节点作为分组ID，也可替换为数字序号）
node_to_group = {node: dsu.find(node) for node in all_nodes}
# 给原表新增分组列：同组两个ID的根节点一致，用任意一列匹配分组即可
so_fake_data['group'] = so_fake_data['id_1'].map(node_to_group)

运行结果 #

最终输出的so_fake_data和预期完全一致：

如果需要用数字作为分组ID，可以增加一步编码映射：

group_mapping = {root: idx+1 for idx, root in enumerate(set(node_to_group.values()))}
so_fake_data['group_num'] = so_fake_data['group'].map(group_mapping)

简化实现（依赖networkx库） #

如果你可以安装第三方库，用networkx实现代码更简洁：

import networkx as nx

# 构建无向图，批量添加所有边
G = nx.from_pandas_edgelist(so_fake_data, source='id_1', target='id_2')
# 获取所有连通分量
components = list(nx.connected_components(G))
# 生成节点到分组编号的映射
node_group_map = {}
for group_id, comp in enumerate(components, 1):
    for node in comp:
        node_group_map[node] = group_id
# 新增分组列
so_fake_data['group'] = so_fake_data['id_1'].map(node_group_map)

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