how to group rows from a pandas data frame that match multiple date-time conditions

Question

I used the code below to create the following data frame:

df = pd.DataFrame({'date':[['2024-02-22 03:42:06.065'],['2024-02-22 03:47:29.152'], ['2024-02-21 19:37:05.142'], ['2024-02-21 19:40:13.851'], ['2024-02-21 19:41:21.388'], ['2024-02-21 19:47:29.828'], ['2024-02-21 19:48:10.684'],['2024-02-22 02:59:36.786']]})

id	Date
AB1234	[2024-02-22 03:42:06.065]
AB1234	[2024-02-22 03:46:29.152]
CD3456	[2024-02-21 19:37:05.142]
CD3456	[2024-02-21 19:40:13.851]
CD3456	[2024-02-21 19:41:21.388]
GH3447	[2024-02-21 19:47:29.828]
GH3447	[2024-02-21 19:48:10.684]
GH3447	[2024-02-22 02:59:36.786]

I need to check if the difference between the dates in 2 (or more) consecutive rows that belong to the same id is <= 5 minutes. (So, if 4 rows belong to the same id, we compare that the diff between row1 and row2, between row2 and row3, and row3 and row4 are <=5 minutes). If yes, group those rows together (below the 'count' column). If not, put them in a different group (as shown in the 'count' column). The purpose is to check if the difference between two consecutive dates that belong to the same id is <= 5 minutes. So, the output should be something like the below:

id	Date	count
AB1234	[2024-02-22 03:42:06.065]	1
AB1234	[2024-02-22 03:46:29.152]	1
CD3456	[2024-02-21 19:37:05.142]	2
CD3456	[2024-02-21 19:40:13.851]	2
CD3456	[2024-02-21 19:41:21.388]	2
GH3447	[2024-02-21 19:47:29.828]	3
GH3447	[2024-02-21 19:48:10.684]	3
GH3447	[2024-02-22 02:59:36.786]	4

Any idea how to do this?

Answer 1

UPDATE

Let's start by defining the data type of your date column (which needs to have the list from each row removed)

df.date = df.date.apply(lambda x: x[0])
df.date = df.date.astype("datetime64[ns]")

which gives us

        id                     date
0   AB1234  2024-02-22 03:42:06.065
1   AB1234  2024-02-22 03:46:29.152
2   CD3456  2024-02-21 19:37:05.142
3   CD3456  2024-02-21 19:40:13.851
4   CD3456  2024-02-21 19:41:21.388

We then use three key methods, df.diff, df.apply, and df.groupby to complete this.

I started by creating a function to run through apply, which either returns the current group (cnt) or creates a new group and updates cnt. Note if your python is <3.8, you cannot use the := operator, instead you can use a standard variable update (see below)

def is_group(df):
    global cnt
    if df["diff"] or not (pd.isna(df["cnt"])):
        return cnt
    else:
    # For Python >=3.8
        return (cnt := cnt + 1)
    # For Python <3.8
        cnt += 1
        return cnt

Then apply this to the grouped dataframes, using diff and the 5 minutes condition to create a mask for where to start groups

cnt = 0
for g, gdf in df.groupby("id"):
    cnt += 1
    gdf["diff"] = gdf.date.diff() <= pd.Timedelta("5 minutes")
    gdf["cnt"] = [cnt] + [None] * (len(gdf) - 1)
    gdf["cnt"] = gdf.apply(is_group, axis=1)
    for idx in gdf.index:
        df.at[idx, "cnt"] = gdf.at[idx, "cnt"]

This gave me

        id                     date cnt
0   AB1234  2024-02-22 03:42:06.065 1.0
1   AB1234  2024-02-22 03:46:29.152 1.0
2   CD3456  2024-02-21 19:37:05.142 2.0
3   CD3456  2024-02-21 19:40:13.851 2.0
4   CD3456  2024-02-21 19:41:21.388 2.0
5   GH3447  2024-02-21 19:47:29.828 3.0
6   GH3447  2024-02-21 19:48:10.684 3.0
7   GH3447  2024-02-22 02:59:36.786 4.0

I tried doing this through a single apply, but couldn't achieve the updating on groups being further than 5 minutes apart.

OLD ANSWER

If you define the date column as a datetime object, which you can do as:

df = pd.DataFrame(
    {
        'date':[
            '2024-02-22 03:42:06.065',
            '2024-02-22 03:47:29.152', 
            '2024-02-21 19:37:05.142', 
            '2024-02-21 19:40:13.851', 
            '2024-02-21 19:41:21.388', 
            '2024-02-21 19:47:29.828', 
            '2024-02-21 19:48:10.684',
            '2024-02-22 02:59:36.786'
        ]
    },
    dtype="datetime64[ns]"
)

You can iterate over the rows to check which rows fit the criteria as below, assuming I understand your logic of any row within 5 minutes of the given row?

for idx, row in df.iterrows():
    df.at[idx, 'count'] = \
        df.loc[abs(df.date - row.date) <= pd.Timedelta("5 minutes"), :] \
            .count() \
                .iloc[1]

For each row, you get all rows (df.iterrows) within 5 minutes (abs(df.date - row.date) <= pd.Timedelta("5 minutes")) for that row compared to all other rows in the dataframe, then take the count which summarises the dataframe to count how many values in each column, and then just simplify to a single value (iloc[1]) instead of a whole dataframe. This is updated at the row index in the count column (df.at[idx, 'count'])