Home > other >  Get mean of N rows with identifier column
Get mean of N rows with identifier column

Time:03-17

I have a pandas dataframe as follows:

data = {'id': ['a1', 'a1', 'a1', 'a2', 'a2', 'a2', 'a2', 'a2', 'a3'], 'val1': [1, 2, 3, 4, 5, 6, 7, 8, 9], 'val2': [10, 20, 30, 40, 50, 60, 70, 80, 90]}
df = pd.DataFrame(data)

which has the following representation:

   id  val1  val2
0  a1     1    10
1  a1     2    20
2  a1     3    30
3  a2     4    40
4  a2     5    50
5  a2     6    60
6  a2     7    70
7  a2     8    80
8  a3     9    90

I want to group the dataframe by id and get the average of val1 and val2 groupped by N rows.

For instance, if N=2, the expected output would be:

   id  val1  val2
0  a1   1.5    15
1  a1     3    30
2  a2   4.5    45
3  a2   6.5    65
4  a2     8    80
5  a3     9    90

As it does the average of each id every 2 elements.

My question is: what is the most efficient way to do this, given that N is provided as a parameter?

CodePudding user response:

Use GroupBy.cumcount with integer division for groups and then aggregate mean:

N = 2
g = df.groupby('id').cumcount() // N
df = df.groupby(['id', g]).mean().droplevel(1).reset_index()
print (df)
   id  val1  val2
0  a1   1.5  15.0
1  a1   3.0  30.0
2  a2   4.5  45.0
3  a2   6.5  65.0
4  a2   8.0  80.0
5  a3   9.0  90.0
  • Related