Python 计算移动平均线

Question

I have a dataset as below:

import pandas as pd
data = {'Category': ['A','A','A','A','A','A','B','B','B','B','B','C','C','C','C','C'],
        'Date' : [1,2,3,4,5,6,1,2,3,4,5,1,2,3,4,5],
       'Count': [1,2,3,4,5,1,2,3,4,5,6,1,2,3,4,6]}
df = pd.DataFrame(data)

I was trying to calculate the average count of every 3 rows excluding the current row grouped by category and the date is from new to old, if there is not enough 3 rows to calculate, it will return 0.

The expected result should be as below. For example: for Category A for Date 1, the average is calculated as average count of dates 2, 3 and 4 of Category A.

Category	Date	Count	Average
A	1	1	3
A	2	2	4
A	3	3	3.333333333
A	4	4	2.666666666
A	5	5	0
A	6	1	0
B	1	2	4
B	2	3	5
B	3	4	0
B	4	5	0
B	5	6	0
C	1	1	3
C	2	2	4.333333333
C	3	3	0
C	4	4	0
C	5	6	0

I was trying to use below, which didn't get expected result

df['average'] = df.groupby(['Category'])['count'].transform(lambda x: x.rolling(3, 1).mean())

Answer 1

You can use rolling in combination with shift and sort_values as follows:

def reverse_roll(df):
    df['Count'] = df.sort_values('Date', ascending=False)['Count'].rolling(3, 3).mean().shift()
    df['Count'] = df['Count'].fillna(0.0)
    df = df.sort_values('Date', ascending=True)
    return df
    
df.groupby('Category').apply(reverse_roll)

Above, rolling(3, 3) is used to force the rolling window to always consider 3 rows and not less. The first 2 rows the result will therefore be NaN which is set to 0 with fillna . To not include the current row in the window, shift is used.

Result:

   Category  Date     Count
0         A     1  3.000000
1         A     2  4.000000
2         A     3  3.333333
3         A     4  0.000000
4         A     5  0.000000
5         A     6  0.000000
6         B     1  4.000000
7         B     2  5.000000
8         B     3  0.000000
9         B     4  0.000000
10        B     5  0.000000
11        C     1  3.000000
12        C     2  4.333333
13        C     3  0.000000
14        C     4  0.000000
15        C     5  0.000000