Pandas : balancing data

g = df.groupby('class')
g.apply(lambda x: x.sample(g.size().min()).reset_index(drop=True))

  class  val
0    c1    1
1    c1    1
2    c2    2
3    c2    2
4    c3    3
5    c3    3

Answers to your follow-up questions

  1. The x in the lambda ends up being a dataframe that is the subset of df represented by the group. Each of these dataframes, one for each group, gets passed through this lambda.
  2. g is the groupby object. I placed it in a named variable because I planned on using it twice. df.groupby('class').size() is an alternative way to do df['class'].value_counts() but since I was going to groupby anyway, I might as well reuse the same groupby, use a size to get the value counts… saves time.
  3. Those numbers are the the index values from df that go with the sampling. I added reset_index(drop=True) to get rid of it.

Leave a Comment