如何在Python中基于person与associate关联生成分组列？

这个问题本质是图论中的连通分量问题：把每个人员看作节点，person和associate的关联关系看作节点间的边，所有能通过直接/间接关联连起来的人员属于同一组。下面提供两种实用的实现方式：

方法一：用NetworkX库（简单快捷） #

NetworkX是Python的图论工具库，能快速处理连通分量问题，适合大多数场景。

代码示例： #

import pandas as pd
import networkx as nx

# 构造示例数据（实际中可替换为pd.read_csv读取你的数据）
data = pd.DataFrame({
    'index': [1,2,3,1,2,3,1,2,3,1,2,3,1,2,3],
    'person': ['name1','name1','name1','name2','name2','name2','name3','name3','name3','name4','name4','name4','name5','name5','name5'],
    'associate': ['-','-','-','name1','name4','name3','-','-','-','-','-','-','-','-','-']
})

# 替换缺失标记为NaN，过滤无效关联行
data['associate'] = data['associate'].replace('-', pd.NA)
valid_edges = data.dropna(subset=['associate'])[['person', 'associate']]

# 创建无向图并添加节点、边
graph = nx.Graph()
all_people = pd.concat([data['person'], data['associate'].dropna()]).unique()
graph.add_nodes_from(all_people)
graph.add_edges_from(valid_edges.values)

# 给每个连通分量分配组ID
group_dict = {}
current_group = 1
for component in nx.connected_components(graph):
    for person in component:
        group_dict[person] = current_group
    current_group += 1

# 映射组ID到原数据，无关联人员留空（或改为单独分组）
data['group'] = data['person'].map(group_dict)

# 可选：给无关联人员分配单独组
# last_group = current_group - 1
# data['group'] = data['group'].fillna(data['person'].rank(method='dense') + last_group)

print(data)

方法二：手动实现并查集（无需额外库） #

如果不能安装第三方库，可以用**并查集（Union-Find）**数据结构手动实现，这是处理连通分量的经典算法。

代码示例： #

import pandas as pd

# 实现并查集类
class UnionFind:
    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):
        if x not in self.parent:
            self.parent[x] = x
        if y not in self.parent:
            self.parent[y] = y
        root_x = self.find(x)
        root_y = self.find(y)
        if root_x != root_y:
            self.parent[root_y] = root_x

# 构造示例数据
data = pd.DataFrame({
    'index': [1,2,3,1,2,3,1,2,3,1,2,3,1,2,3],
    'person': ['name1','name1','name1','name2','name2','name2','name3','name3','name3','name4','name4','name4','name5','name5','name5'],
    'associate': ['-','-','-','name1','name4','name3','-','-','-','-','-','-','-','-','-']
})

# 处理缺失值
data['associate'] = data['associate'].replace('-', pd.NA)
valid_edges = data.dropna(subset=['associate'])[['person', 'associate']]

# 初始化并查集，合并所有关联对
uf = UnionFind()
for _, row in valid_edges.iterrows():
    uf.union(row['person'], row['associate'])

# 给每个人员分配组ID
group_dict = {}
current_group = 1
all_people = pd.concat([data['person'], data['associate'].dropna()]).unique()

# 先处理有关联的人员
for person in all_people:
    if person in uf.parent:
        root = uf.find(person)
        if root not in group_dict:
            group_dict[root] = current_group
            current_group += 1
        group_dict[person] = group_dict[root]

# 给无关联人员分配单独组
for person in data['person'].unique():
    if person not in group_dict:
        group_dict[person] = current_group
        current_group += 1

# 映射到原数据
data['group'] = data['person'].map(group_dict)

print(data)

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