How do I reset the count using the pandas diff() function when a condition from another column is sa-CodePudding

I am trying to count the number of days between dates (cumulatively), (grouped by a column denoted as id), however, I want to reset the counter whenever a condition is satisfied.

Below I have the following dataframe:

     reset_day category       date     id  tdelta
0            N      low 2019-09-04  16876     NaN
1            N      low 2019-09-05  16876     NaN
2            N      low 2019-09-06  16876     NaN
3            N      low 2019-09-07  16876     NaN
4            N      low 2019-09-08  16876     NaN
...        ...      ...        ...    ...     ...
5144         Y   medium 2021-05-23  17612     0.0
5145         Y     high 2021-05-23  23406     0.0
5146         Y     high 2021-05-23  21765     0.0
5147         Y   medium 2021-05-23  19480     0.0
5148         Y   medium 2021-05-23   9066     0.0

Here I want to input values into the column “tdelta”, where the values are currently NaN. This column counts the number of days between the “date” column for each id.

However, I am struggling with trying to reset the count based on the column “reset_day”. If the column value is a “Y”, then the count should start again for that particular id. You can already see a value of 0 in such cases in the tdelta column.

I have previously tried the following on a similar dataframe, by creating another column denoted as test t delta:

example = example.sort_values(by="date")
example['date'] = pd.to_datetime(example['date'], format='%Y/%m/%d')
example['test tdelta'] = example.groupby('id')['date'].diff() / np.timedelta64(1, 'D')
example['test tdelta'] = example['test tdelta'].fillna(0)

However, this just counts the number of days between the dates for each id and fills in the missing values without the resetting I need.

Any ideas on how I can solve this problem??

CodePudding user response：

I think creating an extra grouping column based on the reset day might what you're looking for.

import pandas as pd
import numpy as np

df = pd.DataFrame({'reset_day':['N','N','Y','N','N','Y','Y','Y','Y','Y'],
                   'category':['low','low','low','low','low','medium','high','high','medium','medium'],
                   'date':['2019-09-04','2019-09-05','2019-09-06','2019-09-07','2019-09-08','2021-05-23','2021-05-23','2021-05-23','2021-05-23','2021-05-23'],
                   'id':[16876,16876,16876,16876,16876,17612,23406,21765,19480,9066]
                   })


df['date'] = pd.to_datetime(df['date'], format='%Y/%m/%d')
df = df.sort_values(['id','date'])

#create extra grouping column based on reset day
df['reset_group'] = df['reset_day'].replace({'N':False,'Y':True})
df['reset_group'] = df.groupby('id')['reset_group'].cumsum()

#use extra grouping column when calculating differences
df['tdelta'] = df.groupby(['id','reset_group'])['date'].diff() / np.timedelta64(1, 'D')
df['tdelta'] = df.groupby(['id','reset_group'])['tdelta'].cumsum().fillna(0)

print(df)

  reset_day category       date     id  reset_group  tdelta
9         Y   medium 2021-05-23   9066            1     0.0
0         N      low 2019-09-04  16876            0     0.0
1         N      low 2019-09-05  16876            0     1.0
2         Y      low 2019-09-06  16876            1     0.0
3         N      low 2019-09-07  16876            1     1.0
4         N      low 2019-09-08  16876            1     2.0
5         Y   medium 2021-05-23  17612            1     0.0
8         Y   medium 2021-05-23  19480            1     0.0
7         Y     high 2021-05-23  21765            1     0.0
6         Y     high 2021-05-23  23406            1     0.0