如何使用Pandas归集存在关联关系的元素并分组？

给定如下DataFrame，其中每一行的C1和C2元素存在关联关系，需要将所有存在间接或直接关联的元素归集到同一分组中：

import pandas as pd
df = pd.DataFrame({'C1': ['A', 'B', 'C', 'D', 'R', 'X'], 'C2': ['B', 'C', 'D', 'E', 'S', 'Y']})
# 输出的DataFrame结构：
#    C1  C2
# 0   A   B
# 1   B   C
# 2   C   D
# 3   D   E
# 4   R   S
# 5   X   Y

期望得到的分组结果：

Groups
0  [A, B, C, D, E]
1        [R, S]
2        [X, Y]

这个问题本质是寻找图的连通分量：把每个元素看作节点，每一行的C1-C2看作一条无向边，最终每个连通分量就是一组关联元素。以下提供两种实现方式：

方法1：使用NetworkX库（简洁高效） #

NetworkX是专门处理图结构的Python库，能快速实现连通分量的查找：

import pandas as pd
import networkx as nx

# 初始化输入DataFrame
df = pd.DataFrame({'C1': ['A', 'B', 'C', 'D', 'R', 'X'], 'C2': ['B', 'C', 'D', 'E', 'S', 'Y']})

# 创建无向图并添加所有边
graph = nx.Graph()
graph.add_edges_from(zip(df['C1'], df['C2']))

# 获取所有连通分量，转换为列表格式
connected_groups = [list(component) for component in nx.connected_components(graph)]

# 生成结果DataFrame
result_df = pd.DataFrame({'Groups': connected_groups})
print(result_df)

执行后输出结果与期望完全一致。

方法2：手动实现并查集（无需第三方库） #

如果无法安装第三方库，可以用并查集（Union-Find）算法手动实现分组逻辑，适合环境受限的场景：

import pandas as pd

class UnionFind:
    def __init__(self):
        self.parent = {}
    
    def find(self, node):
        # 查找节点的根节点，路径压缩优化
        if self.parent[node] != node:
            self.parent[node] = self.find(self.parent[node])
        return self.parent[node]
    
    def union(self, node1, node2):
        # 初始化节点（如果未在字典中）
        if node1 not in self.parent:
            self.parent[node1] = node1
        if node2 not in self.parent:
            self.parent[node2] = node2
        # 合并两个节点所在的集合
        root1 = self.find(node1)
        root2 = self.find(node2)
        if root1 != root2:
            self.parent[root2] = root1

# 初始化输入DataFrame
df = pd.DataFrame({'C1': ['A', 'B', 'C', 'D', 'R', 'X'], 'C2': ['B', 'C', 'D', 'E', 'S', 'Y']})

uf = UnionFind()
# 遍历所有关联对，执行合并操作
for c1, c2 in zip(df['C1'], df['C2']):
    uf.union(c1, c2)

# 按根节点分组，整理成结果列表
group_dict = {}
for node in uf.parent:
    root = uf.find(node)
    if root not in group_dict:
        group_dict[root] = []
    group_dict[root].append(node)

# 生成结果DataFrame
result_df = pd.DataFrame({'Groups': list(group_dict.values())})
print(result_df)

两种方法对比 #

• NetworkX方法：代码简洁，开发效率高，适合快速实现；需要额外安装库（pip install networkx）。
• 并查集方法：无需依赖第三方库，运行效率高，适合生产环境或无法安装外部库的场景；代码量稍多。

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