Pandas - General function for checking condition on previous rows-CodePudding

I'm writing a function where I'd like to add a series to a dataframe that is True if n previous rows of another column are negative.

I have this working for a specific number as n, but I don't know how to generalize it into a parameter that could be passed in.

For example, we have a dataframe with a column named Total. The following code will put True in rows in a column titled consecutive_negative_value if there are 6 consecutive rows where Total is less than zero. How can this be generalized to accept any number n instead of just checking six periods?

    df['negative_value'] = df['Total'] < 0
    
    df['consecutive_negative_value'] = ( 
        df['negative_value'] &
        df['negative_value'].shift(1) &
        df['negative_value'].shift(2) &
        df['negative_value'].shift(3) &
        df['negative_value'].shift(4) &
        df['negative_value'].shift(5)
    )

CodePudding user response：

You can use rolling.sum and check if the sum value is equal to the window size:

window = 6
df.Total.lt(0).rolling(window).sum() == window

Example with a window size of 3:

df = pd.DataFrame({'Total': [-1, 2, -2, -2, -3, 2, 3, -1]})

df.Total.lt(0).rolling(3).sum() == 3
0    False
1    False
2    False
3    False
4     True
5    False
6    False
7    False
Name: Total, dtype: bool