如何为Pandas DataFrame中的关联客户组生成唯一ID？

我有一个存储两位客户关系的Pandas DataFrame，数据示例如下。现需为存在关联关系的客户组生成唯一ID，请问该如何实现？
注：实际场景中有数万客户，客户ID为随机数字，示例中的前缀'A'、'B'仅用于演示。

import pandas as pd
  
# initialize list of lists
data = [['A1', 'A1'], 
        ['A1', 'A2'], 
        
        ['A2', 'A1'],
        ['A2', 'A3'],
        
        ['A3', 'A2'],
        ['A3', 'A4'],
        
        ['A4', 'A3'],
        ['A4', 'A5'],
        
        ['B1', 'B1'], 
        ['B1', 'B2'], 
        
        ['B2', 'B1'],
        ['B2', 'B3'],
        
        ['B3', 'B2'],
        ['B3', 'B4'],
        
        ['B4', 'B3'],
        ['B4', 'B5']
        ]
  
# Create the pandas DataFrame
df = pd.DataFrame(data, columns=['cust_1', 'cust_2'])
  
# print dataframe.
df

预期目标：为所有存在关联的客户分配同一个唯一组ID（比如A1-A5为一组，B1-B5为另一组）。

这个问题本质是识别图中的连通分量：将每个客户视为节点，客户间的关联关系视为边，同一个连通分量内的节点即为同一客户组，只需给每个分量分配唯一ID即可。

推荐方法：使用networkx处理连通分量 #

networkx是Python图论工具库，代码直观，能高效处理数万规模的节点数据。

1. 安装依赖（未安装时执行）

pip install networkx

2. 实现代码

import pandas as pd
import networkx as nx

# 初始化数据（与示例一致）
data = [['A1', 'A1'], 
        ['A1', 'A2'], 
        ['A2', 'A1'],
        ['A2', 'A3'],
        ['A3', 'A2'],
        ['A3', 'A4'],
        ['A4', 'A3'],
        ['A4', 'A5'],
        ['B1', 'B1'], 
        ['B1', 'B2'], 
        ['B2', 'B1'],
        ['B2', 'B3'],
        ['B3', 'B2'],
        ['B3', 'B4'],
        ['B4', 'B3'],
        ['B4', 'B5']
        ]

df = pd.DataFrame(data, columns=['cust_1', 'cust_2'])

# 构建无向图：客户间关联是双向的，用无向图处理
G = nx.from_pandas_edgelist(df, 'cust_1', 'cust_2', create_using=nx.Graph())

# 为每个连通分量分配唯一ID
group_map = {}
for group_idx, component in enumerate(nx.connected_components(G), start=1):
    for cust_id in component:
        group_map[cust_id] = group_idx

# 将组ID映射回原DataFrame
df['group_id'] = df['cust_1'].map(group_map)

print(df)

输出结果

cust_1 cust_2  group_id
0      A1     A1         1
1      A1     A2         1
2      A2     A1         1
3      A2     A3         1
4      A3     A2         1
5      A3     A4         1
6      A4     A3         1
7      A4     A5         1
8      B1     B1         2
9      B1     B2         2
10     B2     B1         2
11     B2     B3         2
12     B3     B2         2
13     B3     B4         2
14     B4     B3         2
15     B4     B5         2

方法说明 #

1. 图构建：from_pandas_edgelist直接将DataFrame的关联数据转化为图的边，nx.Graph()确保关联关系双向生效。
2. 连通分量识别：nx.connected_components(G)返回所有互相连通的节点集合，每个集合对应一个客户组。
3. ID映射：遍历每个连通分量，为内部所有客户分配同一个组ID，最后通过map方法将ID绑定到原数据的每一行。

性能适配 #

对于数万客户的规模，networkx的连通分量算法效率足够；若数据量极大，也可替换为scipy.sparse.csgraph.connected_components，但networkx的代码可读性更优。

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