Python DataFrame中关联行的值求和问题

我有一个pandas DataFrame，部分行包含一个ID（ID1）和一个关联ID（ID2）。示例中，a1与a2属于同一关联组（比如对应同一人），而b和c无关联行：

import pandas as pd

test = pd.DataFrame(
    [['a1', 1, 'a2'],
     ['a1', 2, 'a2'],
     ['a1', 3, 'a2'],
     ['a2', 4, 'a1'],
     ['a2', 5, 'a1'],
     ['b',  6,     ],
     ['c',  7,     ]],
    columns=['ID1', 'Value', 'ID2']
)

需要新增一列，计算所有关联行的Value总和，期望输出如下：

我已掌握按ID1分组求和的方法：

test['Group by ID1'] = test.groupby("ID1")["Value"].transform("sum")

但不知道如何结合ID1和ID2实现关联组求和，希望找到非循环的高效方案。

这个问题本质是识别无向图的连通分量：将ID1和ID2看作图的节点，关联关系看作边，同一连通分量内的所有ID共享同一个总和。以下是两种高效实现方案：

方案1：使用networkx库（推荐，高效简洁） #

networkx是专门处理图结构的库，能快速识别连通分量。如果未安装，先执行pip install networkx。

完整代码 #

import pandas as pd
import networkx as nx

test = pd.DataFrame(
    [['a1', 1, 'a2'],
     ['a1', 2, 'a2'],
     ['a1', 3, 'a2'],
     ['a2', 4, 'a1'],
     ['a2', 5, 'a1'],
     ['b',  6,     ],
     ['c',  7,     ]],
    columns=['ID1', 'Value', 'ID2']
)

# 1. 构建无向图
G = nx.Graph()
# 添加所有ID节点（包括ID1和非空的ID2）
all_ids = pd.concat([test['ID1'], test['ID2'].dropna()]).unique()
G.add_nodes_from(all_ids)
# 添加ID1与ID2的关联边
edges = test.dropna(subset=['ID2'])[['ID1', 'ID2']].values
G.add_edges_from(edges)

# 2. 为每个ID分配连通分量标签
component_labels = {}
for idx, component in enumerate(nx.connected_components(G)):
    for node in component:
        component_labels[node] = idx

# 3. 映射标签并计算分组总和
test['component'] = test['ID1'].map(component_labels)
test['Group by ID1 and ID2'] = test.groupby('component')['Value'].transform('sum')

# 输出结果（可移除component列）
print(test.drop('component', axis=1))

执行结果 #

ID1  Value  ID2  Group by ID1 and ID2
0  a1      1   a2                    15
1  a1      2   a2                    15
2  a1      3   a2                    15
3  a2      4   a1                    15
4  a2      5   a1                    15
5   b      6  NaN                     6
6   c      7  NaN                     7

方案2：无额外依赖实现（纯pandas） #

如果不想引入外部库，可以通过迭代合并关联ID的方式生成分组标签，适合小规模数据：

完整代码 #

import pandas as pd

test = pd.DataFrame(
    [['a1', 1, 'a2'],
     ['a1', 2, 'a2'],
     ['a1', 3, 'a2'],
     ['a2', 4, 'a1'],
     ['a2', 5, 'a1'],
     ['b',  6,     ],
     ['c',  7,     ]],
    columns=['ID1', 'Value', 'ID2']
)

def get_component_mapping(df):
    # 初始化每个ID的分组为自身
    all_ids = pd.concat([df['ID1'], df['ID2'].dropna()]).unique()
    mapping = {id_val: id_val for id_val in all_ids}
    
    # 迭代合并关联ID的分组
    changed = True
    while changed:
        changed = False
        for _, row in df.dropna(subset=['ID2']).iterrows():
            id1, id2 = row['ID1'], row['ID2']
            root1 = mapping[id1]
            root2 = mapping[id2]
            if root1 != root2:
                # 将root2的所有映射替换为root1
                for k in mapping:
                    if mapping[k] == root2:
                        mapping[k] = root1
                changed = True
    return mapping

# 生成分组映射并计算总和
component_mapping = get_component_mapping(test)
test['component'] = test['ID1'].map(component_mapping)
test['Group by ID1 and ID2'] = test.groupby('component')['Value'].transform('sum')

# 输出结果
print(test.drop('component', axis=1))

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