An efficient way to add a new column to Pandas dataframe based on sub-groups of another column

user_3556304 0
Viewed 49

I have a dataframe with various columns as below:

import numpy as np
import pandas as pd

d = {'X': [1.1, 1.2, 1.3, 2.1, 2.2, 2.3],
     'Group Index': [1, 1, 1, 2, 2, 2],
     'Score': [1, 2, 3, 10, 20, 30]}
df = pd.DataFrame(data=d)
print(df)


 X    Group Index  Score
1.1        1         1
1.2        1         2
1.3        1         3
2.1        2        10
2.2        2        20
2.3        2        30

I use "apply" and "lambda" to create a new column to populate the range (e.g., max-min) for each row that has identical "Group Index"

df.loc[:,'Score Range'] = df.apply (lambda row: np.ptp(df[(df['Group Index']==row['Group Index'])]['Score']), axis=1)
print(df)

 X   Group Index  Score  Score Range
1.1       1         1         2
1.2       1         2         2
1.3       1         3         2
2.1       2        10        20
2.2       2        20        20
2.3       2        30        20

The "lambda" and "apply" works, but it is too slow. Is there another better way to speed-up generating the above-noted new column (Score Range)?

pandas dataframe numpy python
2 Answers

So let us try transform + groupby

df['Score Range'] = df.groupby('Group Index')['Score'].transform(np.ptp)
user_79645270

You can also try lazy groupby:

groups = df.groupby('Group Index')['Score']
df['Score Range'] = groups.transform('max') - groups.transform('min')
user_42384080
Related