Search code examples
pandaspandas-groupbytransformation

Transform dataframe based on three first rows in pandas

I have a data frame like this (but much bigger) and I am trying to use transform to get the max based ONLY on the first 3 rows of each group.

     df10 = pd.DataFrame({
      'Price': [1,2,3,4,5,10,20,30,40,50],
      'Stock': ['AAPL', 'AAPL', 'AAPL', 'AAPL', 'AAPL', 'IBM','IBM','IBM','IBM','IBM']
     })

This syntax works for the entire column

df10['max_top_3']=df10.groupby("Stock").Price.transform('max')

But I want the 'max_top_3' column to show 3 and 30 respectively for AAPL and IBM >> which is the max number of the first 3 entries in that column

I tried something like this and it gave an error

df10['max_top_3']=df10.groupby("Stock").Price.head(3).transform('max')
Solution

  • You can chain the head in transform with a lambda:

    df10.groupby("Stock").Price.transform(lambda x: x.head(3).max())

    0     3
1     3
2     3
3     3
4     3
5    30
6    30
7    30
8    30
9    30
Name: Price, dtype: int64