'Find' sequences of size 3 in a dataframe column using pandas

user_12060206 0
Viewed 60

I have a datetime dataset that contains a column where all 'days off' are identified as 1s and the rest are 0s. I am trying to create a new column that identifies extended periods off, say 3 days or more. I need to either select those sequences of 1s greater than 3 in a row (eg: 01110) or discard single (010) and doubles 1s (0110).

An example dataset follows:

import pandas as pd
import numpy as np

df = pd.DataFrame({'days_off': [1,0,1,1,1,0,0,1,1,0,1,0,1,1,1,1,0,1,1,1]})

df['extended_off'] = np.NaN

I've tried a convoluted for loop to look at all the conditions, it appears to be working except at the first two rows, which i cant iterate over with my solution (below)

Is there a better way???, since you really should avoid looping over a df.

for i in range(2, len(df)):     
    if ((df.loc[i-1, 'days_off'] == 0) and (df.loc[i+1, 'days_off'] == 0)): # single holiday (wednesday)
        df.loc[i, 'extended_off'] = 0
    elif ((df.loc[i-1, 'days_off'] == 0) and (df.loc[i+2, 'days_off'] == 0)): # normal weekend (no prior)
        df.loc[i, 'extended_off'] = 0
    elif ((df.loc[i-2, 'days_off'] == 0) and (df.loc[i+1, 'days_off'] == 0)): # normal weekend (no  following)
        df.loc[i, 'extended_off'] = 0
    elif df.loc[i, 'days_off'] == 0: # normal working day
        df.loc[i, 'extended_off'] = 0
    else:
        df.loc[i, 'extended_off'] = 1 # 3 or more days_off in a row

Thanks

pandas dataframe for-loop python
1 Answers

Maybe this could be useful

df["new"] = df['days_off'].ne(df['days_off'].shift()).cumsum()
df["counts"] = df.groupby("new")['new'].transform('size')
((df.counts.ge(3)) & (df.days_off == 1)).astype(int)
user_144023800
Related