如何基于关联ID对Pandas DataFrame多列分组并高效聚合属性

我有一个包含多列ID的Pandas DataFrame，部分ID存在缺失或重复情况，需要基于这些关联ID对DataFrame进行分组，将指定属性以列表形式保留，其他属性保留有效单一值。

示例数据 #

import pandas as pd
import numpy as np

# 初始化数据
data = {
    'first_id': ['id_1', 'id_2', np.nan, 'id_3','id_2'],
    'second_id': ['aaa', np.nan, 'aaa', 'bbb', 'ccc'],
    'third_id': ['db_01', 'db_02', np.nan, np.nan, np.nan],
    'sources': [1,1,2,1,3]
}

# 创建DataFrame
df = pd.DataFrame(data)

原始DataFrame：

first_id second_id third_id  sources
0     id_1       aaa    db_01        1
1     id_2       NaN    db_02        1
2      NaN       aaa      NaN        2
3     id_3       bbb      NaN        1
4     id_2       ccc      NaN        3

期望输出 #

最终需要得到按关联ID分组后的结果，同一关联组的sources合并为去重列表，其他ID列保留该组内的有效非空值：

result_data = {
    'first_id': ['id_1', 'id_2','id_3'],
    'second_id': ['aaa', 'ccc', 'bbb'],
    'third_id': ['db_01', 'db_02', np.nan],
    'sources': [[1,2],[1,3],[1]]
}

pd.DataFrame(result_data)

输出：

first_id second_id third_id  sources
0     id_1       aaa    db_01  [1, 2]
1     id_2       ccc    db_02  [1, 3]
2     id_3       bbb      NaN     [1]

现有实现问题 #

我编写了分组函数，但当前输出中id_1的sources仅得到[1]，未关联到行2的sources=2，且函数需要分步处理，希望得到更简洁的解决方案：

def get_element_from_pandas(col):
    """从列表中取第一个元素，用于Pandas列处理"""
    if col is np.nan or type(col) == float:
        return col
    else:
        if len(col) != 0:
            return col[0]
        else:
            return col


def group_dataframe(df: pd.DataFrame, index_to_group: list, group_columns: list, keep_in_list: list):
    """分组处理DataFrame"""
    df = df.groupby(index_to_group, as_index=False)[group_columns].agg(lambda x: [*dict.fromkeys(x)])

    for col in group_columns:
        # 移除列表中的nan
        df[col] = df[col].apply(lambda x: [i for i in x if str(i) != "nan"])
        # 空列表替换为nan
        df[col] = df[col].apply(lambda x: np.nan if len(x) == 0 else x)

        if col not in keep_in_list:
            df[col] = df[col].apply(get_element_from_pandas)

    return df


# 调用函数后的输出
group_dataframe(df, ['first_id'], ['second_id', 'third_id', 'sources'], ['sources'])

输出结果：

first_id second_id third_id  sources
0     id_1       aaa    db_01     [1]
1     id_2       ccc    db_02  [1, 3]
2     id_3       bbb      NaN     [1]

简洁解决方案 #

核心问题是要先识别所有关联ID的连通组（比如id_1和aaa属于同一组，因此行0和行2要合并），可以用图论工具networkx来实现：

完整代码 #

import pandas as pd
import numpy as np
import networkx as nx

# 1. 构建关联图：将每一行的非空ID视为节点，同一行的节点互相连接
G = nx.Graph()
for _, row in df.iterrows():
    # 提取当前行所有非空ID
    ids = [val for val in row[['first_id', 'second_id', 'third_id']] if pd.notna(val)]
    # 同一行的节点两两连边
    if len(ids) > 1:
        for i in range(len(ids)):
            for j in range(i+1, len(ids)):
                G.add_edge(ids[i], ids[j])
    elif len(ids) == 1:
        G.add_node(ids[0])

# 2. 获取每个节点的连通分量标签
component_labels = {}
for idx, component in enumerate(nx.connected_components(G)):
    for node in component:
        component_labels[node] = idx

# 3. 给原DataFrame添加组标签：取当前行第一个非空ID对应的组标签
def get_group_label(row):
    for col in ['first_id', 'second_id', 'third_id']:
        val = row[col]
        if pd.notna(val):
            return component_labels[val]
    return np.nan

df['group_id'] = df.apply(get_group_label, axis=1)

# 4. 按组ID聚合
agg_funcs = {
    'first_id': lambda x: next((val for val in x if pd.notna(val)), np.nan),
    'second_id': lambda x: next((val for val in x if pd.notna(val)), np.nan),
    'third_id': lambda x: next((val for val in x if pd.notna(val)), np.nan),
    'sources': lambda x: list(pd.unique(x))
}

result = df.groupby('group_id', as_index=False).agg(agg_funcs).drop(columns='group_id')

# 调整列顺序与期望输出一致
result = result[['first_id', 'second_id', 'third_id', 'sources']]
print(result)

输出结果 #

first_id second_id third_id  sources
0     id_1       aaa    db_01  [1, 2]
1     id_2       ccc    db_02  [1, 3]
2     id_3       bbb      NaN     [1]

代码说明 #

• 构建关联图：把每一行的非空ID作为图的节点，同一行的节点之间建立连接，这样所有关联的ID会处于同一个连通分量中。
• 分配组标签：给每个连通分量分配唯一ID，确保关联的ID属于同一组。
• 分组聚合：按组ID聚合，每个ID列取组内第一个非空值，sources列取去重后的列表。

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