Drop sorted row based on value count column

Question

My dataframe looks like this:

   year   id    
0  2019   x1
1  2012   x1
2  2017   x1
3  2013   x1
4  2018   x2
5  2012   x2
6  2013   x2

I want to filter my whole dataframe such that if there are more than 3 observations per id, the observation with the lowest year should be dropped.

In this case, the 1th row should be dropped.

   year   id    
0  2019   x1
1  2017   x1
2  2013   x1
3  2018   x2
4  2012   x2
5  2013   x2

Answer 1

Use DataFrame.sort_values with GroupBy.head :

df = df.sort_values(['id','year'], ascending=[True, False]).groupby('id').head(3)
print (df)
   year  id
0  2019  x1
2  2017  x1
3  2013  x1
4  2018  x2
6  2013  x2
5  2012  x2

If order should be same add DataFrame.sort_index :

df = df.sort_values(['id','year'], ascending=[True, False]).groupby('id').head(3).sort_index()
print (df)
   year  id
0  2019  x1
2  2017  x1
3  2013  x1
4  2018  x2
5  2012  x2
6  2013  x2

Answer 2

Using GroupBy.nlargest :

df = df.groupby('id')['year'].nlargest(3).reset_index().drop(columns='level_1')

   id  year
0  x1  2019
1  x1  2017
2  x1  2013
3  x2  2018
4  x2  2013
5  x2  2012

Make sure that year has an int dtype:

df['year'] = df['year'].astype(int)

Answer 3

What about using a for loop for solving this problem (I love for Loops):

id_unique = df.id.unique()

df_new = pd.DataFrame(columns = df.columns)

for i in id_unique:
    df_new = pd.concat([df_new, df[df['id'] == i ].sort_values(['year'], ascending= [False]).head(3)], axis=0)