Map two dataframes, count events where timestamps in second dataframe are within the date-time ranges of the first dataframe

Viewed 74

I have a dataframe df that records the movement of people from one location to another. I have a second dataframe df2 that records events at certain timeframes. For each entry in df, I would like to count the number of events, per event type, if the event timestamp occured between start_ts and end_ts given that the person_id's match.

import pandas as pd
import numpy as np

current_time = '2022-05-05 17:00'

df = pd.DataFrame({"person_id": ["A", "B", "A", "C", "A", "C"],
                   "location_id": [1, 5, 2, 7, 3, 8],
                   "start_ts": pd.to_datetime(["2022-05-05 00:00", "2022-05-05 00:00", "2022-05-05 05:00", "2022-05-05 00:00", "2022-05-05 13:00", "2022-05-05 11:00"]),
                   "end_ts": pd.to_datetime(["2022-05-05 02:00", "2022-05-05 03:00", "2022-05-05 10:00", "2022-05-05 04:00", "2022-05-05 16:00", "2022-05-05  12:00"]),
})

df2 = pd.DataFrame({"person_id": ["A", "A", "A", "A", "A", "A", "A", "A", "A", "B", "B", "B", "C", "C", "C", "C", "C", "C"],
                    "timestamp": pd.to_datetime(["2022-05-05 01:00", "2022-05-05 01:10", "2022-05-05 01:30", "2022-05-05 06:00",
                                             "2022-05-05 07:00", "2022-05-05 08:00", "2022-05-05 13:00", "2022-05-05 14:00",
                                             "2022-05-05 15:00", "2022-05-05 01:00", "2022-05-05 01:30", "2022-05-05 02:00",
                                             "2022-05-05 01:00", "2022-05-05 02:00", "2022-05-05 03:00", "2022-05-05 11:10",
                                             "2022-05-05 11:20", "2022-05-05 11:30"]),
                    "event": ["1", "2", "3", "1", "2", "2", "3", "1", "1", "2", "3", "3", "1", "1", "3", "2", "3", "1"],
})

df looks as follows:

print(df.to_string())
>>
  person_id  location_id            start_ts              end_ts
0         A            1 2022-05-05 00:00:00 2022-05-05 02:00:00
1         B            5 2022-05-05 00:00:00 2022-05-05 03:00:00
2         A            2 2022-05-05 05:00:00 2022-05-05 10:00:00
3         C            7 2022-05-05 00:00:00 2022-05-05 04:00:00
4         A            3 2022-05-05 13:00:00 2022-05-05 16:00:00
5         C            8 2022-05-05 11:00:00 2022-05-05 12:00:00

df2 looks as follows:

print(df2.to_string())
>>
   person_id           timestamp event
0          A 2022-05-05 01:00:00     1
1          A 2022-05-05 01:10:00     2
2          A 2022-05-05 01:30:00     3
3          A 2022-05-05 06:00:00     1
4          A 2022-05-05 07:00:00     2
5          A 2022-05-05 08:00:00     2
6          A 2022-05-05 13:00:00     3
7          A 2022-05-05 14:00:00     1
8          A 2022-05-05 15:00:00     1
9          B 2022-05-05 01:00:00     2
10         B 2022-05-05 01:30:00     3
11         B 2022-05-05 02:00:00     3
12         C 2022-05-05 01:00:00     1
13         C 2022-05-05 02:00:00     1
14         C 2022-05-05 03:00:00     3
15         C 2022-05-05 11:10:00     2
16         C 2022-05-05 11:20:00     3
17         C 2022-05-05 11:30:00     1

My attempt at doing the pivot type count is as follows:

df2 = df2.sample(frac=1)

for i in df.index:
    for j in np.unique(df2["event"]):
        df.loc[i, "count_event_" + j] = len(df2.loc[(df2["timestamp"].between(df.loc[i, "start_ts"], df.loc[i, "end_ts"], inclusive = "both")) & 
                                                (df2["event"] == j) &
                                                (df2["person_id"] == df.loc[i, "person_id"])])

My desired output is:

print(df.to_string())
>>
  person_id  location_id            start_ts              end_ts  count_event_1  count_event_2      count_event_3
0         A            1 2022-05-05 00:00:00 2022-05-05 02:00:00            1.0            1.0            1.0
1         B            5 2022-05-05 00:00:00 2022-05-05 03:00:00            0.0            1.0            2.0
2         A            2 2022-05-05 05:00:00 2022-05-05 10:00:00            1.0            2.0            0.0
3         C            7 2022-05-05 00:00:00 2022-05-05 04:00:00            2.0            0.0            1.0
4         A            3 2022-05-05 13:00:00 2022-05-05 16:00:00            2.0            0.0            1.0
5         C            8 2022-05-05 11:00:00 2022-05-05 12:00:00            1.0            1.0            1.0

I would like to get rid of the nested for loops, and instead, achieve the same result using a group by or pivot. How can I achieve this efficiently, especially considering that the actual dataframes are a few million entries each.

1 Answers

Here is one way to do it (I've left intermediate print-outs to help better understand what's happening at each step):

# Setup
df["start_ts"] = pd.to_datetime(df["start_ts"], format="%Y-%m-%d %H:%M:%s")
df["end_ts"] = pd.to_datetime(df["end_ts"], format="%Y-%m-%d %H:%M:%s")
df2["timestamp"] = pd.to_datetime(df2["timestamp"], format="%Y-%m-%d %H:%M:%s")
# Find matching intervals
for idx in df["person_id"].unique():
    df2.loc[df2["person_id"] == idx, "interval"] = df2.loc[
        df2["person_id"] == idx, "timestamp"
    ].map(
        lambda x: df.loc[
            (df["person_id"] == idx) & (df["start_ts"] <= x) & (x <= df["end_ts"]),
            ["start_ts", "end_ts"],
        ].index[0]
    )

print(df2.head())
# Output
  person_id           timestamp event  interval
0         A 2022-05-05 01:00:00     1       0.0
1         A 2022-05-05 01:10:00     2       0.0
2         A 2022-05-05 01:30:00     3       0.0
3         A 2022-05-05 06:00:00     1       2.0
4         A 2022-05-05 07:00:00     2       2.0
# Count number of events
df2 = (
    df2.assign(interval=lambda x: x["interval"].astype(int))
    .assign(count=1)
    .groupby(["person_id", "interval", "event"])
    .agg({"count": sum})
    .reset_index(drop=False)
    .pivot(index=["person_id", "interval"], columns="event")
    .reset_index(drop=False)
)
df2.columns = df2.columns.droplevel()
df2.columns = [
    "person_id",
    "interval",
    "count_event_1",
    "count_event_2",
    "count_event_3",
]

print(df2)
# Output
  person_id  interval    1    2    3
0         A         0  1.0  1.0  1.0
1         A         2  1.0  2.0  NaN
2         A         4  2.0  NaN  1.0
3         B         1  NaN  1.0  2.0
4         C         3  2.0  NaN  1.0
5         C         5  1.0  1.0  1.0
# Final dataframe
df = df.reset_index(drop=False).rename(columns={"index": "interval"})
df = (
    pd.merge(left=df, right=df2, on=["person_id", "interval"])
    .fillna(0)
    .astype(int, errors="ignore")
)

print(df)
# Output as expected
   interval person_id  location_id            start_ts              end_ts  count_event_1  count_event_2  count_event_3
0         0         A            1 2022-05-05 00:00:00 2022-05-05 02:00:00              1              1              1
1         1         B            5 2022-05-05 00:00:00 2022-05-05 03:00:00              0              1              2
2         2         A            2 2022-05-05 05:00:00 2022-05-05 10:00:00              1              2              0
3         3         C            7 2022-05-05 00:00:00 2022-05-05 04:00:00              2              0              1
4         4         A            3 2022-05-05 13:00:00 2022-05-05 16:00:00              2              0              1
5         5         C            8 2022-05-05 11:00:00 2022-05-05 12:00:00              1              1              1
Related