CODEWITHSUNDEEP
pythonpandasgroup-bygreatest-n-per-group
user6566438
user6566438

Reputation: 229

Extracting group observations from pandas dataframe

I have a pandas dataframe. I want to extract a certain number of observations from each sub group of the dataframe and put them into a new dataframe. For example, let's assume we have the following dataframe:

         Var1    Var2
    0     1      1.2
    1     2      1.3
    2     2      1.4
    3     1      1.5
    4     1      1.6
    5     2      1.7
    6     1      1.8
    7     1      1.9
    8     2      2.0
    9     1      2.1
   10     2      2.2
   11     1      2.3

I want to sort it by var1 first:

         Var1    Var2
    0     1      1.2
    1     1      1.5
    2     1      1.6
    3     1      1.8
    4     1      1.9
    5     1      2.1
    6     1      2.3
    7     2      1.3
    8     2      1.4
    9     2      1.7
   10     2      2.0
   11     2      2.2

and then keep the first two observations of each group and put them to a new dataframe:

         Var1    Var2
    0     1      1.2
    1     1      1.5
    2     2      1.3
    3     2      1.4

I know how to use group by, but it is not clear to me how to perform the second step. Thank you very much for the help.

Upvotes: 1

Views: 194

Answers (1)

jezrael
jezrael

Reputation: 862841

Use sort_values with groupby and head:

df = df.sort_values('Var1').groupby('Var1').head(2).reset_index(drop=True)
print (df)
   Var1  Var2
0     1   1.2
1     1   1.5
2     2   1.3
3     2   1.4

df = df.groupby('Var1').head(2).sort_values('Var1').reset_index(drop=True)
print (df)
   Var1  Var2
0     1   1.2
1     1   1.5
2     2   1.3
3     2   1.4

Another solution with iloc:

df = df.groupby('Var1')['Var2']
       .apply(lambda x: x.iloc[:2])
       .reset_index(level=1, drop=True)
       .reset_index()
print (df)
   Var1  Var2
0     1   1.2
1     1   1.5
2     2   1.3
3     2   1.4

Note:

For older version of pandas change sort_values to sort, but rather toupgrade to last version.

Upvotes: 2

Related Questions