How to compute mean for different size subsets within pandas dataframe?

user_7712225 0
Viewed 55

compute mean of particular column for each unique subset of rows in pandas dataframe. In following example each subset is till 1 appears in column "Flag" i.e. (54+34+78+91+29)/5 = 57.2 and (81+44+61)/3 = 62.0

Currently unable to compute rolling subset of different sizes based on particular column condition


>>> import pandas as pd
>>> df = pd.DataFrame({"Indx": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], "Units": [54, 34, 78, 91, 29, 81, 44, 61, 73, 19], "Flag": [0, 0, 0, 0, 1, 0, 0, 1, 0, 1]})
>>> df
   Indx  Units  Flag
0     1     54     0
1     2     34     0
2     3     78     0
3     4     91     0
4     5     29     1
5     6     81     0
6     7     44     0
7     8     61     1
8     9     73     0
9    10     19     1

# DESIRED OUTPUT
>>> df
   Indx  Units  Flag   avg
0     1     54     0  57.2
1     2     34     0  57.2
2     3     78     0  57.2
3     4     91     0  57.2
4     5     29     1  57.2
5     6     81     0  62.0
6     7     44     0  62.0
7     8     61     1  62.0
8     9     73     0  46.0
9    10     19     0  46.0
pandas mean subset dataframe
2 Answers

Create the group key by using cumsum then transform 

df['Units'].groupby(df.Flag.iloc[::-1].cumsum()).transform('mean')
0    57.2
1    57.2
2    57.2
3    57.2
4    57.2
5    62.0
6    62.0
7    62.0
8    46.0
9    46.0
Name: Units, dtype: float64

#df['new']=df['Units'].groupby(df.Flag.iloc[::-1].cumsum()).transform('mean')
user_79645270

The shortest solution (I think) is:

df['avg'] = df.groupby(df.Flag[::-1].cumsum()).Units.transform('mean')

You don't even need to use iloc, as df.Flag[::-1] retrieves Flag column in reversed order.

user_73884770
Related