在Pandas中对共享REG的ID进行分组编号

需要在Pandas DataFrame中，将所有共享同一REG编号的ID归为同一分组，并新增GROUP列标记分组编号。

原始数据：

ID       REG
15      01497
15      01493
19      01706
19      01706-A
78      05710
78      01738
143     01626
143     01634
144     01626
144     01644
207     05255
207     01638
209     05255
209     03143

期望结果：

ID     REG          GROUP
15     01497        1
15     01493        1
19     01706        2
19     01706-A      2
78     05710        3
78     01738        3
143    01626        4
143    01634        4
144    01626        4
144    01644        4
207    05255        5
207    01638        5
209    05255        5
209    03143        5

注：需求并非传统的求和或计数聚合操作，而是基于ID与REG的关联关系进行分组。

这个问题本质是找连通分量：ID和REG构成关联网络，共享REG的ID属于同一连通分量，需要给每个连通分量分配唯一分组编号。可以用networkx库实现，步骤如下：

• 构建ID与REG的关联图：将ID和REG都作为图的节点，每条数据记录作为连接ID和REG的边
• 识别图中的所有连通分量
• 为每个连通分量分配GROUP编号，并映射回原DataFrame

代码实现 #

import pandas as pd
import networkx as nx

# 加载原始数据（可替换为从文件读取，比如pd.read_csv）
data = [
    [15, "01497"], [15, "01493"],
    [19, "01706"], [19, "01706-A"],
    [78, "05710"], [78, "01738"],
    [143, "01626"], [143, "01634"],
    [144, "01626"], [144, "01644"],
    [207, "05255"], [207, "01638"],
    [209, "05255"], [209, "03143"]
]
df = pd.DataFrame(data, columns=["ID", "REG"])

# 构建关联图
G = nx.Graph()
for _, row in df.iterrows():
    G.add_edge(row["ID"], row["REG"])

# 为连通分量分配分组编号
group_mapping = {}
current_group = 1
for component in nx.connected_components(G):
    for node in component:
        group_mapping[node] = current_group
    current_group += 1

# 新增GROUP列
df["GROUP"] = df["ID"].map(group_mapping)

# 输出结果
print(df)

说明 #

• 此方法确保只要两个ID共享任意REG，就会被归到同一分组（比如ID143和144共享REG01626，同属GROUP4；ID207和209共享REG05255，同属GROUP5）
• 若未安装networkx，执行pip install networkx即可完成安装

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