如何对含双ID列的DataFrame分组并汇总所有关联数据？

我有一个包含两个ID字段的DataFrame，结构如下：

我希望将所有关联的ID归为一组：ID 5关联6和7，而3因关联7也需归入该组。预期结果如下：

尝试过简单的group_by方法但无效，因为需要对另一ID列进行拼接操作，阻碍了按第二ID分组；曾考虑用grepl匹配列表分组，但仍未找到可行方法。

这个问题本质是图的连通分量识别问题：ID1和ID2的关联关系可看作图中的边，所有相互连通的ID属于同一组。以下是两种可行实现方式：

方法一：使用networkx库处理连通分量 #

import pandas as pd
import networkx as nx

# 构造原始DataFrame
df = pd.DataFrame({
    'ID 1': [5,5,6,3],
    'ID 2': [6,7,7,7],
    'Name': ['Apple1', 'Applee', 'Applye', 'Apple2']
})

# 创建无向图并添加节点与边
G = nx.Graph()
all_ids = pd.concat([df['ID 1'], df['ID 2']]).unique()
G.add_nodes_from(all_ids)
edges = list(df[['ID 1', 'ID 2']].itertuples(index=False))
G.add_edges_from(edges)

# 生成ID到组标识的映射（以5为目标组的标识）
id_to_group = {}
for comp in nx.connected_components(G):
    group_id = 5 if 5 in comp else min(comp)
    for id_val in comp:
        id_to_group[id_val] = group_id

# 添加组标识并聚合数据
df['group_id'] = df['ID 1'].map(id_to_group)
result = df.groupby('group_id').agg(
    **{
        'ID 1': lambda x: x.iloc[0],
        'ID 2': lambda x: ', '.join(map(str, set(pd.concat([df['ID 1'], df['ID 2']])[pd.concat([df['ID 1'], df['ID 2']]).map(id_to_group) == x.name] - {x.name}))),
        'Name': lambda x: ', '.join(x)
    }
).reset_index(drop=True)

方法二：使用并查集（Union-Find）算法实现 #

无需依赖第三方图库，用基础算法实现连通分量识别：

import pandas as pd

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):
        x_root = self.find(x)
        y_root = self.find(y)
        if x_root != y_root:
            self.parent[y_root] = x_root

# 构造原始DataFrame
df = pd.DataFrame({
    'ID 1': [5,5,6,3],
    'ID 2': [6,7,7,7],
    'Name': ['Apple1', 'Applee', 'Applye', 'Apple2']
})

# 初始化并查集并合并关联ID
uf = UnionFind()
all_ids = pd.concat([df['ID 1'], df['ID 2']]).unique()
for id_val in all_ids:
    uf.parent[id_val] = id_val

for _, row in df.iterrows():
    uf.union(row['ID 1'], row['ID 2'])

# 生成ID到组标识的映射（以5为目标组的标识）
root_5 = uf.find(5)
id_to_group = {}
for id_val in all_ids:
    id_to_group[id_val] = 5 if uf.find(id_val) == root_5 else uf.find(id_val)

# 添加组标识并聚合数据
df['group_id'] = df['ID 1'].map(id_to_group)
result = df.groupby('group_id').agg(
    **{
        'ID 1': lambda x: x.iloc[0],
        'ID 2': lambda x: ', '.join(map(str, set(pd.concat([df['ID 1'], df['ID 2']])[pd.concat([df['ID 1'], df['ID 2']]).map(id_to_group) == x.name] - {x.name}))),
        'Name': lambda x: ', '.join(x)
    }
).reset_index(drop=True)

两种方法最终都会得到符合预期的聚合结果，解决普通groupby无法处理的跨ID关联分组问题。

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