如何用Pandas或NetworkX实现基于Name/Country的关联聚类分组？

这个需求本质是寻找无向图的连通分量：每个ID是节点，若两个ID的Name或Country相同，则节点间存在边，最终连通的节点组成一个簇。NetworkX是更合适的工具，它原生支持连通分量计算，代码简洁高效；Pandas也能实现，但需要手动处理迭代合并，相对繁琐。

一、NetworkX实现（推荐） #

通过引入中间节点（Name和Country的标识），可以高效建立节点间的连通关系，无需两两连接所有同Name/Country的ID，大幅减少边的数量。

import pandas as pd
import networkx as nx

# 初始化DataFrame
df = pd.DataFrame(
    {'Name':['Smith', 'Brown', 'Smith', 'Miller'], 
     'Country': ['US', 'GB', 'DE','US']}, 
    index=[0,1,2,3]
)
df.index.name = 'ID'

# 创建无向图
G = nx.Graph()
# 添加所有ID作为节点
G.add_nodes_from(df.index)

# 为每个Name建立ID与Name节点的连接（用前缀区分Name和Country节点）
for name, group in df.groupby('Name'):
    G.add_edges_from((id, f"Name_{name}") for id in group.index)

# 为每个Country建立ID与Country节点的连接
for country, group in df.groupby('Country'):
    G.add_edges_from((id, f"Country_{country}") for id in group.index)

# 提取连通分量，过滤掉中间的Name/Country节点，只保留ID
clusters = []
for component in nx.connected_components(G):
    id_cluster = [node for node in component if isinstance(node, int)]
    if id_cluster:
        clusters.append(id_cluster)

# 转换为要求的字典格式
result = {f'cluster{i}': sorted(cluster) for i, cluster in enumerate(clusters)}
print(result)

输出结果：

{'cluster0': [0, 2, 3], 'cluster1': [1]}

二、Pandas实现（备选） #

通过迭代合并分组，不断扩展簇的范围，直到簇不再变化。这种方法适合小数据集，大数据集下效率较低。

import pandas as pd

# 初始化DataFrame
df = pd.DataFrame(
    {'Name':['Smith', 'Brown', 'Smith', 'Miller'], 
     'Country': ['US', 'GB', 'DE','US']}, 
    index=[0,1,2,3]
)
df.index.name = 'ID'

# 初始簇：按Name分组得到的ID集合
clusters = [set(group.index) for _, group in df.groupby('Name')]

# 迭代合并簇，直到没有新的合并发生
changed = True
while changed:
    changed = False
    new_clusters = []
    while clusters:
        current_cluster = clusters.pop()
        # 获取当前簇所有ID对应的Country
        related_countries = df.loc[current_cluster, 'Country'].unique()
        # 找到所有属于这些Country的ID
        related_ids = set(df[df['Country'].isin(related_countries)].index)
        
        # 如果当前簇可以扩展，合并并标记变化
        if current_cluster != related_ids:
            current_cluster = current_cluster.union(related_ids)
            changed = True
            # 合并new_clusters中与当前簇重叠的簇
            to_merge = []
            for idx, c in enumerate(new_clusters):
                if c & current_cluster:
                    to_merge.append(idx)
            # 倒序合并避免索引错乱
            for idx in reversed(to_merge):
                current_cluster = current_cluster.union(new_clusters.pop(idx))
        new_clusters.append(current_cluster)
    clusters = new_clusters

# 去重并转换为字典格式
clusters = list(set(frozenset(c) for c in clusters))
result = {f'cluster{i}': sorted(cluster) for i, cluster in enumerate(clusters)}
print(result)

输出结果与NetworkX方法一致。

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