Pandas groupby time gaps

I have a dataframe called df_merged which outputs like below:

Here is the info on this df:

<class 'pandas.core.frame.DataFrame'>
RangeIndex: 154752 entries, 0 to 154751
Data columns (total 5 columns):

     Column         Non-Null Count   Dtype                  
 0   count_date     154752 non-null  datetime64[ns]
 1   time_gap       0 non-null       category      
 2   device         154752 non-null  object        
 3   count          154752 non-null  int32         
 4   average_speed  154752 non-null  int32         
dtypes: category(1), datetime64[ns](1), int32(2), object(1)
memory usage: 3.7+ MB

time_gap is a category because I replaced it with a pd.cut() function. I don't know if it's best to change the dtype here to groupby.

I would like to groupby time_gap, knowing that the average_speed has to be weighted by count with this function:

def average_speed_mean(x):
    try: 
        return np.average(x["average_speed"], weights=x["count"])
    except ZeroDivisionError:
        return 0

I'm trying to have basically the same dataframe grouped like this:

I tried this:

df_merged = df_merged.groupby("time_gap").agg({("count", sum), ("average_speed", average_speed_mean)})

But it doesn't seem to work out and I have no idea how I could solve this.

Thank you in advance for your help.

Does the following give you the result you want:

groups = ["count_date", "time_gap", "device"]
result = (
    df_merged
    .assign(average_speed=df_merged["average_speed"] * df_merged["count"])
    .groupby(groups).agg({"count": "sum", "average_speed": "sum"})
    .assign(average_speed=lambda df: df["average_speed"].div(df["count"]).fillna(0))
    .reset_index()
)

• groups is the list of columns I think you want to group over - adjust it, if that's not the case.
• Multiply the average_speed by the count.
• Now group over the groups and sum up the counts and average_speeds per group.
• Then divide the summed up average_speed by the total count and fill then NaN with 0.
• Finally reset the index to get the group columns back.

Result for the small sample in the question is:

   count_date        time_gap device  count  average_speed
0  2018-12-05 0 days 00:00:00  CAT17      0            0.0
1  2018-12-05 0 days 01:00:00  CAT17      0            0.0
2  2023-05-04 0 days 22:00:00  CAT17      0            0.0
3  2023-05-04 0 days 23:00:00  CAT17      8           16.0

Thank you, it's working very well and I find your solution pretty neat to understand it step by step.

Here is the output I get: