How to balance a dataset in `dplyr` using `sample_n` automatically to the size of the smallest class?

Question

I have a dataset like:

df <- tibble(
  id = 1:18,
  class = rep(c(rep(1,3),rep(2,2),3),3),
  var_a = rep(c("a","b"),9)
)

# A tibble: 18 x 3
      id cluster var_a
   <int>   <dbl> <chr>
 1     1       1 a    
 2     2       1 b    
 3     3       1 a    
 4     4       2 b    
 5     5       2 a    
 6     6       3 b    
 7     7       1 a    
 8     8       1 b    
 9     9       1 a    
10    10       2 b    
11    11       2 a    
12    12       3 b    
13    13       1 a    
14    14       1 b    
15    15       1 a    
16    16       2 b    
17    17       2 a    
18    18       3 b 

That dataset contains a number of observations in several classes. The classes are not balanced. In the sample above we can see, that only 3 observations are of class 3, while there are 6 observations of class 2 and 9 observations of class 1.

Now I want to automatically balance that dataset so that all classes are of the same size. So I want a dataset of 9 rows, 3 rows in each class. I can use the sample_n function from dplyr to do such a sampling.

I achieved to do so by first calculating the smallest class size..

min_length <- as.numeric(df %>% 
  group_by(class) %>% 
  summarise(n = n()) %>% 
  ungroup() %>% 
  summarise(min = min(n)))

..and then apply the sample_n function:

set.seed(1)
df %>% group_by(cluster) %>% sample_n(min_length)

# A tibble: 9 x 3
# Groups:   cluster [3]
     id cluster var_a
  <int>   <dbl> <chr>
1    15       1 a    
2     7       1 a    
3    13       1 a    
4     4       2 b    
5     5       2 a    
6    17       2 a    
7    18       3 b    
8     6       3 b    
9    12       3 b    

I wondered If it's possible to do that (calculating the smallest class size and then sampling) in one go?

Answer 1

You can do it in one step, but it is cheating a little:

set.seed(42)
df %>%
  group_by(class) %>%
  sample_n(min(table(df$class))) %>%
  ungroup()
# # A tibble: 9 x 3
#      id class var_a
#   <int> <dbl> <chr>
# 1     1     1 a    
# 2     8     1 b    
# 3    15     1 a    
# 4     4     2 b    
# 5     5     2 a    
# 6    11     2 a    
# 7    12     3 b    
# 8    18     3 b    
# 9     6     3 b    

I say "cheating" because normally you would not want to reference df$ from within the pipe. However, because they property we're looking for is of the whole frame but the table function only sees one group at a time, we need to side-step that a little.

One could do

df %>%
  mutate(mn = min(table(class))) %>%
  group_by(class) %>%
  sample_n(mn[1]) %>%
  ungroup()
# # A tibble: 9 x 4
#      id class var_a    mn
#   <int> <dbl> <chr> <int>
# 1    14     1 b         3
# 2    13     1 a         3
# 3     7     1 a         3
# 4     4     2 b         3
# 5    16     2 b         3
# 6     5     2 a         3
# 7    12     3 b         3
# 8    18     3 b         3
# 9     6     3 b         3

Though I don't think that that is any more elegant/readable.