如何基于Pandas DataFrame两列中的唯一ID进行分组？

问题背景 #

现有长度为n的ID列表ids，以及包含列A、B、C的DataFrame df。已知ids中的每个ID仅会出现在A或B列中的一列（不会同时出现在两列），需要按这些ID对df进行分组，最终得到以ID为键、对应C列数据列表的结果。

示例 #

输入：

import pandas as pd
ids = [1, 2, 3, 4, 5]
df = pd.DataFrame({
    "A" : [1, 6, 3, 7, 5, 1],
    "B" : [8, 2, 9, 3, 10, 11],
    "C" : ["data1", "data2", "data3", "data4", "data5", "data6"]
})

期望输出：

{
    1: ["data1", "data6"],
    2: ["data2"],
    3: ["data3", "data4"],
    4: [],
    5: ["data5"]
}

之前尝试df.groupby(["A,B"])，但该写法会将A、B列作为联合分组键，无法达到需求，以下是几种简洁高效的解决方案：

解决方案 #

方法一：用np.where快速生成分组键（高效推荐） #

利用numpy的向量化操作生成分组ID，比逐行apply效率高，适合大规模数据：

import numpy as np

# 标记A列ID是否在目标列表中
mask = df['A'].isin(ids)
# 生成分组ID：A列符合则取A，否则取B
df['group_id'] = np.where(mask, df['A'], df['B'])

# 过滤出目标ID的行，分组聚合C列并转字典
grouped = df[df['group_id'].isin(ids)].groupby('group_id')['C'].agg(list).to_dict()

# 补全无数据的ID，默认空列表
result = {id_val: grouped.get(id_val, []) for id_val in ids}

方法二：重塑数据后分组（代码简洁） #

通过melt将A、B列转为单列，再进行分组，逻辑清晰：

# 重塑数据：将A、B列合并为group_id列，保留C列
melted_df = df.melt(id_vars=['C'], value_vars=['A', 'B'], value_name='group_id')

# 过滤目标ID并分组聚合
grouped = melted_df[melted_df['group_id'].isin(ids)].groupby('group_id')['C'].agg(list).to_dict()

# 补全无数据的ID
result = {id_val: grouped.get(id_val, []) for id_val in ids}

方法三：逐行提取分组ID（直观易懂） #

用apply逐行判断并提取分组ID，适合中小规模数据，逻辑最直观：

# 逐行提取属于ids的ID（A或B）
df['group_id'] = df.apply(lambda row: row['A'] if row['A'] in ids else row['B'], axis=1)

# 过滤+分组+转字典
grouped = df[df['group_id'].isin(ids)].groupby('group_id')['C'].apply(list).to_dict()

# 补全无数据的ID
for id_val in ids:
    if id_val not in grouped:
        grouped[id_val] = []
result = grouped

注意事项 #

• 由于题目保证每个ID仅出现在A或B列之一，无需处理同一行A/B均在ids中的冲突情况。
• 数据量较大时优先选择方法一，np.where是向量化操作，远快于逐行apply；方法二的melt效率也不错，但会生成临时DataFrame。

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