как ускорить этот расчет?

Question

У меня есть набор данных, который выглядит следующим образом:

| ..... userId ................. | ..cahtroomID ... . | ..msg_index_in_chat .. | ..time_difference_between_msg .. | | 1234567891222222 | sdfbsjkfdsdklf ... | .............. 1 ................. | ...... 0 часов 0 минут .................... | | 9876543112252141 | sdfbsjkfdsdklf ... | ............... 2 ................ | ...... 0 часов 4 минут .................... | | 2374623982398939 | quweioqewiieio | ............... 1 ................ | ...... 0 часов 0 минут .. .................. | | 1234567891222222 | quweioqewiieio | ............... 2 ................ | ...... 0 часов 4 минуты .. .................. | | 2374623982398939 | quweioqewiieio | .................... 3 ........... | ...... 1 час 0 минут .. .................. |

Мне нужно вычислить среднее время между сообщениями в каждой комнате и назначить значение, которое я получил для каждой строки. Для этого я написал:

 df['avg_time'] = 0
    for room in set(df.roomId):
        table = df[['msg_index_in_chat', 'time_difference_between_msg']][df.roomId == room]
        if len(table) > 1:
            avg_time = []
            times = table.time_difference_between_msg.tolist()
            avg_time = sum(times[1:], timedelta(0))/len(times[1:])
        elif len(table) ==1:
            avg_time = timedelta(hours = 0)
        df.loc[df.roomId == room,('avg_time')] = avg_time

проблема в том, что этот код выполняется за лот времени. Можете ли вы предложить более быстрый способ сделать этот расчет?

Спасибо!

Answer 1

Использование GroupBy.transform с пользовательской лямбда-функцией:

f = lambda times: sum(times.iloc[1:], pd.Timedelta(0))/len(times.iloc[1:]) if len(times) > 1 else pd.Timedelta(0)
df['avg_time'] = df.groupby('cahtroomID')['time_difference_between_msg'].transform(f)
print (df)
             userId      cahtroomID  msg_index_in_chat  \
0  1234567891222222  sdfbsjkfdsdklf                  1   
1  9876543112252141  sdfbsjkfdsdklf                  2   
2  2374623982398939  quweioqewiieio                  1   
3  1234567891222222  quweioqewiieio                  2   
4  2374623982398939  quweioqewiieio                  3   

  time_difference_between_msg avg_time  
0                    00:00:00 00:04:00  
1                    00:04:00 00:04:00  
2                    00:00:00 00:32:00  
3                    00:04:00 00:32:00  
4                    01:00:00 00:32:00