Python pandas: apply a function to dataframe.rolling()

Yi Fang picture Yi Fang · Apr 15, 2018 · Viewed 8.6k times · Source

I have this dataframe:

In[1]df = pd.DataFrame([[1,2,3,4,5],[6,7,8,9,10],[11,12,13,14,15],[16,17,18,19,20],[21,22,23,24,25]])
In[2]df
Out[2]: 
    0   1   2   3   4
0   1   2   3   4   5
1   6   7   8   9  10
2  11  12  13  14  15
3  16  17  18  19  20
4  21  22  23  24  25

I need to achieve this:

  1. for every rows in my dataframe,
  2. if 2 or more values within any 3 consecutive cells is greater than 10,
  3. then the last of that 3 cells should be marked as True.

The resulting dataframe df1 should be same size with True of False in it based on the above stated criteria:

In[3]df1
Out[3]: 
    0   1      2      3      4
0 NaN NaN  False  False  False
1 NaN NaN  False  False  False
2 NaN NaN   True   True   True
3 NaN NaN   True   True   True
4 NaN NaN   True   True   True
  • df1.iloc[0,1] is NaN bacause in that cell, only two numbers were given but needed atleast 3 numbers to do the test.
  • df1.iloc[1,3] is False since none in [7,8,9] is greater than 10
  • df1.iloc[3,4] is True since 2 or more in [18,19,20] is greater than 10

I figured dataframe.rolling.apply() with a function might be the solution, but how exactly?

Answer

penguin2048 picture penguin2048 · Apr 15, 2018

You are right that using rolling() is the way to go. However, you must keep in mind since rolling() replaces the value at end of the window with the new value, so you can not just mark the window with True you will also get False whenever the condition is not applicable

Here is the code that uses your sample dataframe and performs the desired transformation:

df = pd.DataFrame([[1,2,3,4,5],[6,7,8,9,10],[11,12,13,14,15],[16,17,18,19,20],[21,22,23,24,25]])

now, defining a function that takes a window as an argument and returns whether the condition is satisfied

def fun(x):
    num = 0
    for i in x:
        num += 1 if i > 10 else 0
    return 1 if num >= 2 else -1

I have hardcoded the threshold as 10. So if in any window the numbers of values greater than 10 are greater than or equal to 2 than the last value is replaced by 1 (denoting True), else it is replaced by -1(denoting False).

If you want to keep threshold parameters as variables, then have a look at this answer to pass them as arguments.

Now applying the function on rolling window, using window size as 3, axis 1 and additionally if you don't want NaN then you can also set min_periods to 1 in the arguments.

df.rolling(3, axis=1).apply(fun)

produces the output as

  0   1    2    3    4
0 NaN NaN -1.0 -1.0 -1.0
1 NaN NaN -1.0 -1.0 -1.0
2 NaN NaN  1.0  1.0  1.0
3 NaN NaN  1.0  1.0  1.0
4 NaN NaN  1.0  1.0  1.0