Pandas Groupby循环下快速为DataFrame分配行的优化方案

我有两个DataFrame：

import pandas as pd

df = pd.DataFrame({'A':['German Shepherd','Border Collie','Golden Retriever','Beagle','Daschund']})
df = df.T
df.columns = df.iloc[0]
df = df.drop(df.index[0])

此时df的结构为：

第二个DataFrame：

df2 = pd.DataFrame({'ID':['A','A','A','B','C','C','C','C','C'],
                   'Breed':['German Shepherd','Beagle','Dashung','Border Collie',
                           'German Shepherd','Border Collie','Golden Retriever','Beagle','Daschund']})

df2的结构为：

找出df2中每个犬种所属的ID，据此更新df：若某ID下存在该犬种，对应位置标记为1，否则为0。

以下代码可正常运行，但处理大型数据集（50万行，生成4000×30000的矩阵）时速度极慢，耗时数小时：

import numpy as np

dogs_grouped = df2.groupby('ID')
missing_dogs = []
vals = [np.nan for i in df.columns]
for group_name, df_group in dogs_grouped:
    print(f'Cluster: {group_name}')
    cluster_dogs = sorted(list(set(df_group['Breed'].to_list())))
    cluster_dogs = [i for i in cluster_dogs if i in all_dogs]  # all_dogs为df的列名集合
    weird_dogs = [i for i in cluster_dogs if i not in all_dogs]
    missing_dogs.append(weird_dogs)
    df = df.append(pd.Series(vals, index=df.columns, name=group_name))
    df.loc[group_name][cluster_dogs] = 1
df = df.fillna(0)

利用Pandas矢量化操作替代循环，大幅提升处理速度：

步骤1：生成ID-犬种存在矩阵 #

先对df2去重（同一ID下同一犬种仅需标记一次），再用透视表生成二进制矩阵：

# 去重，避免重复统计同一ID下的同一犬种
df2_unique = df2.drop_duplicates(subset=['ID', 'Breed'])
# 生成透视表，存在标记为1，不存在标记为0
pivot_df = df2_unique.pivot_table(index='ID', columns='Breed', aggfunc='size', fill_value=0)

步骤2：对齐列并合并结果 #

将透视表的列与df的列对齐，仅保留目标犬种，补充缺失值：

# 仅保留df中指定的犬种列，缺失列填充0
result = pivot_df.reindex(columns=df.columns, fill_value=0)
# 合并原df的初始行（若有），并填充剩余NaN为0
result = result.combine_first(df).fillna(0)

步骤3：收集无效犬种（可选） #

如果需要统计df2中不在df列中的犬种（如示例中的Dashung）：

all_dogs = set(df.columns)
weird_dogs = df2[~df2['Breed'].isin(all_dogs)]['Breed'].unique().tolist()

最终结果 #

处理后得到的矩阵如下：

效率说明 #

该方案完全依赖Pandas的底层矢量化运算，避免了循环中逐行append和loc赋值的低效操作，处理大型数据集时速度可提升数个数量级，4000×30000的矩阵可在分钟级完成。

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