Groupby 和命名聚合 | 優化 Pandas 中的數據幀生成

Question

我在 Pandas 中有一個包含一些列的數據框，如下所示：

data = {
    'CODIGO_SINIESTRO': [10476434, 10476434, 4482524, 4482524, 4486110],
    'CONDICION': ['PASAJERO', 'CONDUCTOR', 'MOTOCICLISTA', 'CICLISTA', 'PEATON'],
    'EDAD': [62.0, 29.0, 26.0, 47.0, 33.0],
    'SEXO': ['MASCULINO', 'FEMENINO', 'FEMENINO', 'MASCULINO', 'FEMENINO']
}

df = pd.DataFrame(data)

輸出：

    CODIGO_SINIESTRO    CONDICION       EDAD    SEXO
0   10476434            PASAJERO        62.0    MASCULINO
1   10476434            CONDUCTOR       29.0    MASCULINO
2   4482524             MOTOCICLISTA    26.0    MASCULINO
3   4482524             CICLISTA        47.0    MASCULINO
4   4486110             PEATON          33.0    FEMENINO

因此，我想創建另一個按'CODIGO_SINIESTRO'列分組的數據'CODIGO_SINIESTRO' ，並且我想要以下列如結果：

• 'CODIGO_SINIESTRO' ：行的 ID。
• 'PROMEDIO_EDAD' ：此列將存儲 edad 均值。
• 'CANTIDAD_HOMBRES' ：此列將存儲基於'SEXO'列的男性計數。
• 'CANTIDAD_HOMBRES' ：此列將存儲基於'SEXO'列的女性計數。

最后，我想要五個額外的列，命名為等於'CONDICION'列的四個可能值，如果值存在，則該值將存儲 1，否則將存儲 0。

所以，我編寫了這個解決方案並按預期工作，但是我的數據集中有很多行（150k+）並且解決方案很慢（5 分鍾）。 這是我的代碼：

df_final = df.groupby(['CODIGO_SINIESTRO']).agg(
    CANTIDAD_HOMBRES=pd.NamedAgg(column='SEXO', aggfunc=lambda x: (x=='MASCULINO').sum()),
    CANTIDAD_MUJERES=pd.NamedAgg(column='SEXO', aggfunc=lambda x: (x=='FEMENINO').sum()),
    PROMEDIO_EDAD=pd.NamedAgg(column='EDAD', aggfunc=np.mean),
    MOTOCICLISTA=pd.NamedAgg(column='CONDICION', aggfunc=lambda x: (x=='MOTOCICLISTA').any().astype(int)),
    CONDUCTOR=pd.NamedAgg(column='CONDICION', aggfunc=lambda x: (x=='CONDUCTOR').any().astype(int)),
    PEATON=pd.NamedAgg(column='CONDICION', aggfunc=lambda x: (x=='PEATON').any().astype(int)),
    CICLISTA=pd.NamedAgg(column='CONDICION', aggfunc=lambda x: (x=='CICLISTA').any().astype(int)),
    PASAJERO=pd.NamedAgg(column='CONDICION', aggfunc=lambda x: (x=='PASAJERO').any().astype(int))
).reset_index()

輸出：

    CODIGO_SINIESTRO    CANTIDAD_HOMBRES    CANTIDAD_MUJERES    PROMEDIO_EDAD ...    
                                                    
 0    4482524                  1                      1               36.5  
 1    4486110                  0                      1               33.0  
 2    10476434                 1                      1               45.5


... MOTOCICLISTA    CONDUCTOR   PEATON  CICLISTA    PASAJERO
        1               0         0        1           0
        0               0         1        0           0
        0               1         0        0           1

我該如何優化這個解決方案？，還有其他方法可以解決這個問題嗎？

謝謝你。

Answer 1

使用矢量化方法進行預聚合應該效率更高（事實證明它快了 100 倍）：

df['PROMEDIO_EDAD']= df.groupby('CODIGO_SINIESTRO')['EDAD'].transform(np.mean)
df['CANTIDAD_HOMBRES'] = np.where(df['SEXO'] == 'MASCULINO', 1, 0)
df['CANTIDAD_MUJERES'] = np.where(df['SEXO'] == 'FEMENINO', 1, 0)
for col in df['CONDICION'].unique():
    df[col] = np.where(df['CONDICION'] == col, 1, 0)
df = df.groupby(['CODIGO_SINIESTRO', 'PROMEDIO_EDAD']).sum().reset_index().drop('EDAD', axis=1)
df.iloc[:,2:] = (df.iloc[:,2:] > 0).astype(int)
df
Out[1]: 
   CODIGO_SINIESTRO  PROMEDIO_EDAD  CANTIDAD_HOMBRES  CANTIDAD_MUJERES  \
0           4482524           36.5                 1                 1   
1           4486110           33.0                 0                 1   
2          10476434           45.5                 1                 1   

   PASAJERO  CONDUCTOR  MOTOCICLISTA  CICLISTA  PEATON  
0         0          0             1         1       0  
1         0          0             0         0       1  
2         1          1             0         0       0