仅用Pandas实现两列间至少匹配一次的分组（禁用Networkx）

现有如下数据表：

+------+---------------+
| id   |   co_id       |
+------+---------------+
| 123  | abc           |   
| 123  | bcd           |
| 234  |abc            |  
| 123  |def            |
| 678  |urt            |
| 678  |urt            |
| 124  |xyz            |
| 600  |urt            |
| 700  |def            |
+------+---------------+

其中id与co_id存在多对多、一对多或多对一关系（至少存在一次匹配）。目标是仅使用Pandas将所有关联的id和co_id归为同一分组，得到如下格式的输出：

+------+---------------+------+
| id   |   co_id       |cat   |
+------+---------------+------+
|123   |abc            |group1|
|123   |bcd            |group1|
|234   |abc            |group1|
|123   |def            |group1|
|700   |def            |group1|
|678   |urt            |group2|
|678   |urt            |group2|
|600   |urt            |group2|
|124   |xyz            |group3|
+------+---------------+------+

用户尝试的初始代码如下：

import pandas as pd

df = pd.read_csv(r'diretory\sample.csv')

df_id = df.groupby('id').agg({'id': lambda x: list(x)}).reset_index()
df_coid = df.groupby('co_id').agg({'co_id': lambda x: list(x)}).reset_index()

这个问题本质是寻找连通分量：将id和co_id视为图中的节点，每一行的(id, co_id)是连接两个节点的边，所有连通的节点属于同一分组。我们可以用**并查集（Union-Find）**算法配合Pandas实现，无需额外第三方库：

步骤1：实现并查集核心函数 #

def find(u, parent):
    # 查找节点u的根节点，带路径压缩优化
    while parent[u] != u:
        parent[u] = parent[parent[u]]
        u = parent[u]
    return u

def union(u, v, parent):
    # 合并u和v所在的集合
    u_root = find(u, parent)
    v_root = find(v, parent)
    if u_root != v_root:
        parent[v_root] = u_root

步骤2：处理数据并生成分组 #

import pandas as pd

# 读取数据，注意路径替换为你的实际路径
df = pd.read_csv(r'diretory\sample.csv')

# 清理数据：去除co_id字段的首尾空格，统一id为字符串类型（与co_id类型匹配）
df['co_id'] = df['co_id'].str.strip()
df['id'] = df['id'].astype(str)

# 收集所有唯一节点：id和co_id的去重值
all_nodes = pd.concat([df['id'], df['co_id']]).unique()

# 初始化并查集：每个节点的父节点初始化为自身
parent = {node: node for node in all_nodes}

# 遍历每一行，将当前id和co_id所在的集合合并
for _, row in df.iterrows():
    union(row['id'], row['co_id'], parent)

# 为每个连通分量分配组名
root_to_group = {}
current_group = 1
for node in all_nodes:
    root = find(node, parent)
    if root not in root_to_group:
        root_to_group[root] = f'group{current_group}'
        current_group += 1

# 将组名映射回原DataFrame
df['cat'] = df['id'].apply(lambda x: root_to_group[find(x, parent)])

# 按组排序，与示例输出格式对齐
df = df.sort_values('cat').reset_index(drop=True)

# 查看结果
print(df)

代码说明 #

1. 数据预处理：统一节点类型、清理空格，避免因类型或格式不一致导致的匹配失败。
2. 并查集初始化：为每个节点创建初始父节点。
3. 集合合并：遍历每一行数据，将关联的id和co_id合并到同一集合。
4. 组名分配：为每个连通分量的根节点分配唯一组名，再映射回原表。
5. 排序对齐：按组名排序，使输出与示例格式一致。

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