如何基于双条件对Pandas DataFrame行进行分组？

问题背景 #

我有一个包含Phase、Approach、info1、info2列的Pandas DataFrame，需要基于以下两个条件对行进行分组：

• 条件1：具有相同Phase和Approach的行归为一组
• 条件2：Phase不同但Approach、info1、info2完全相同的行也归为一组

示例数据：

预期分组结果：

• 0和1（满足条件1）
• 2和3（满足条件2）
• 4和5（满足条件2）
• 6（无匹配）
• 7和8（满足条件1）

之前尝试分开按两个条件分组，但无法让它们同时生效，该如何实现？

解决方案：基于连通分量的分组 #

这个问题本质是将满足任一条件的行视为连通节点，最终的分组就是图中的连通分量。用networkx库可以高效实现这个逻辑，步骤如下：

代码实现 #

import pandas as pd
import networkx as nx

# 创建示例DataFrame
data = {
    'Phase': ['A1', 'A1', 'A1', 'A2', 'A2', 'A3', 'A3', 'B', 'B'],
    'Approach': ['EB', 'EB', 'WB', 'WB', 'EB', 'EB', 'WB', 'NB', 'NB'],
    'info1': ['J', 'N', 'L', 'L', 'L', 'L', 'S', 'T', 'K'],
    'info2': ['K', 'V', 'M', 'M', 'M', 'M', 'H', 'K', 'T']
}
df = pd.DataFrame(data)

# 初始化图结构
G = nx.Graph()
G.add_nodes_from(df.index)

# 添加条件1的边：相同Phase+Approach的行两两相连
for (phase, approach), group in df.groupby(['Phase', 'Approach']):
    if len(group) > 1:
        nodes = group.index.tolist()
        for i in range(len(nodes)):
            for j in range(i+1, len(nodes)):
                G.add_edge(nodes[i], nodes[j])

# 添加条件2的边：Approach+info1+info2相同且Phase不同的行两两相连
for (approach, info1, info2), group in df.groupby(['Approach', 'info1', 'info2']):
    if len(group) > 1 and group['Phase'].nunique() > 1:
        nodes = group.index.tolist()
        for i in range(len(nodes)):
            for j in range(i+1, len(nodes)):
                G.add_edge(nodes[i], nodes[j])

# 计算连通分量，为每行分配分组ID
connected_components = list(nx.connected_components(G))
group_id_map = {}
for idx, component in enumerate(connected_components):
    for node in component:
        group_id_map[node] = idx

# 将分组ID写入原DataFrame
df['group_id'] = df.index.map(group_id_map)

# 输出结果
print(df)

输出结果 #

Phase Approach info1 info2  group_id
0    A1       EB     J     K         0
1    A1       EB     N     V         0
2    A1       WB     L     M         1
3    A2       WB     L     M         1
4    A2       EB     L     M         2
5    A3       EB     L     M         2
6    A3       WB     S     H         3
7     B       NB     T     K         4
8     B       NB     K     T         4

结果说明 #

• group_id 0：包含行0、1（满足条件1）
• group_id 1：包含行2、3（满足条件2）
• group_id 2：包含行4、5（满足条件2）
• group_id 3：仅行6（无匹配条件）
• group_id 4：包含行7、8（满足条件1）

完全符合预期分组要求。

替代方案：无额外库实现 #

如果不想引入networkx，可以用迭代合并的方式实现：

import pandas as pd

df = pd.DataFrame(data)
# 初始化每个行的组ID为自身索引
df['group_id'] = df.index

# 处理条件1：合并相同Phase+Approach的组
for (phase, approach), group in df.groupby(['Phase', 'Approach']):
    if len(group) > 1:
        base_id = group['group_id'].iloc[0]
        df.loc[group.index, 'group_id'] = base_id

# 处理条件2：合并相同Approach+info1+info2且Phase不同的组
for (approach, info1, info2), group in df.groupby(['Approach', 'info1', 'info2']):
    if len(group) > 1 and group['Phase'].nunique() > 1:
        base_id = group['group_id'].min()
        df.loc[group.index, 'group_id'] = base_id

# 处理连锁合并（比如A和B合并，B和C合并，需同步A的组ID）
changed = True
while changed:
    changed = False
    for current_id in df['group_id'].unique():
        current_rows = df[df['group_id'] == current_id]
        # 找到所有关联的组ID
        cond1_related = df[(df['Phase'].isin(current_rows['Phase'])) & 
                           (df['Approach'].isin(current_rows['Approach']))]['group_id'].unique()
        cond2_related = df[(df['Approach'].isin(current_rows['Approach'])) & 
                           (df['info1'].isin(current_rows['info1'])) & 
                           (df['info2'].isin(current_rows['info2']))]['group_id'].unique()
        all_related = set(cond1_related) | set(cond2_related)
        if len(all_related) > 1:
            new_id = min(all_related)
            if df['group_id'].isin(all_related).any():
                df.loc[df['group_id'].isin(all_related), 'group_id'] = new_id
                changed = True

print(df)

这个方案无需额外依赖，但逻辑相对复杂，需要处理连锁合并的场景。

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