Pandas：基于多列关联分组为记录分配相同聚类ID

需要为大型Pandas DataFrame新增cluster列，规则是：存在直接关联（同一列值相同）或间接关联（通过中间记录的列值传递关联）的记录，分配相同的聚类ID。数据集规模超3000万条，需高效实现。

原始数据 #

表格展示 #

代码生成 #

import pandas as pd
df = pd.DataFrame({
  'A': [111,111,111,222,222,222,333,444], 
  'B': [222,222,111,222,333,444,555,666],
  'C': [111,222,555,444,111,333,777,777]
})

需求示例 #

前6条记录通过A、B、C列的关联形成一个聚类（比如记录1与2因A=111关联，记录2与4因B=222关联，记录1与5因C=111关联），后2条通过C=777关联形成另一个聚类，期望结果如下：

期望结果表格 #

高效解决方案（适配3000万级数据） #

针对超大规模数据集，**并查集（Disjoint Set Union, DSU）**是最优选择——时间复杂度接近线性，远优于图遍历类方法。以下是具体实现：

1. 实现并查集类 #

class DSU:
    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

2. 构建关联并合并集合 #

将每条记录的索引与各列值关联，合并同一列值对应的所有索引：

dsu = DSU()
# 初始化：每个行索引作为独立节点
for idx in df.index:
    dsu.parent[idx] = idx

# 遍历所有列，合并同一列值下的所有行索引
for col in ['A', 'B', 'C']:
    # 按列值分组，获取每组的行索引列表
    groups = df.groupby(col).groups
    for _, indices in groups.items():
        # 取组内第一个索引为基准，合并组内其他所有索引
        base_idx = indices[0]
        for idx in indices[1:]:
            dsu.union(base_idx, idx)

3. 生成连续聚类ID #

将每个索引的根节点映射为连续的聚类ID：

# 获取每个索引对应的根节点
df['root'] = df.index.map(dsu.find)
# 将根节点映射为从1开始的连续ID
root_to_cluster = {root: i+1 for i, root in enumerate(df['root'].unique())}
df['cluster'] = df['root'].map(root_to_cluster)
# 可选：删除中间辅助列
df.drop('root', axis=1, inplace=True)

4. 大数据优化建议 #

• 分块处理：若内存不足，可按列值分块读取数据，逐块合并关联关系，最后统一映射聚类ID
• 替换数据结构：用numpy数组替代字典存储父节点，进一步提升运算速度
• 跳过重复合并：可在合并前判断节点是否已在同一集合，减少不必要操作（并查集的路径压缩已大幅降低重复计算）

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