Pandas实现任意关联公共ID的连通DataFrame行合并

这个需求本质是求解无向图的连通分量，常规低效实现一般是多层循环做重复匹配，时间复杂度高，用带路径压缩、按秩合并优化的并查集（DSU）实现可以做到接近线性时间复杂度，1100行数据可以毫秒级出结果，后续行数增长也不会有性能问题。

测试样例说明 #

输入为3列存储长整型ID的DataFrame，同一ID可出现在多行多列，需要将任意层级通过公共ID关联的所有ID归为同一组，样例如下：

# 输入DataFrame（A/B/C为第一行的ID值，非列名）
A   B   C
X   Y   Z
D   E   F
T   U   V
C   D   E
E   N   Z
AA  BB  CC
HH  CC  U

# 预期输出（组内ID顺序不影响）
组1: A B C D E F N Z X Y
组2: T U V HH CC AA BB

实现步骤 #

• 初始化带路径压缩、按秩合并优化的并查集结构，支持接近常数时间的集合合并、根节点查询操作
• 逐行遍历DataFrame，将每行的3个ID合并到同一集合——同一行的三个ID属于直接关联关系
• 遍历所有出现过的ID，按照所属集合的根节点分组，同一组内的ID即为所求的连通集合

可直接运行的代码 #

import pandas as pd

class DSU:
    def __init__(self):
        self.parent = dict()
        self.rank = dict()

    def find(self, x):
        if x not in self.parent:
            self.parent[x] = x
            self.rank[x] = 1
        # 路径压缩优化
        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:
            return
        # 按秩合并优化
        if self.rank[x_root] < self.rank[y_root]:
            self.parent[x_root] = y_root
        else:
            self.parent[y_root] = x_root
            if self.rank[x_root] == self.rank[y_root]:
                self.rank[x_root] += 1

# 替换为你自己的DataFrame即可
df = pd.DataFrame(
    [
        ["A", "B", "C"],
        ["X", "Y", "Z"],
        ["D", "E", "F"],
        ["T", "U", "V"],
        ["C", "D", "E"],
        ["E", "N", "Z"],
        ["AA", "BB", "CC"],
        ["HH", "CC", "U"],
    ],
    columns=["col1", "col2", "col3"]
)

dsu = DSU()
# 注意：如果你的DataFrame已有自定义列名，把下方col1/col2/col3替换成你实际的列名即可
for _, row in df.iterrows():
    id1, id2, id3 = row["col1"], row["col2"], row["col3"]
    dsu.union(id1, id2)
    dsu.union(id2, id3)

# 按连通分量分组得到结果
res = dict()
for id_val in dsu.parent:
    root = dsu.find(id_val)
    res.setdefault(root, []).append(id_val)

# 打印输出结果
for group in res.values():
    print(" ".join(group))

性能说明 #

该实现时间复杂度接近O(n)（n为所有ID的总数量），1100行数据处理耗时在10毫秒以内，即使后续行数增长到10万级别，也不会出现耗时过长的问题。

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