dplyr sample_n 按组，每组具有唯一的大小参数

Question

I am trying to draw a stratified sample from a data set for which a variable exists that indicates how large the sample size per group should be.

library(dplyr)
# example data 
df <- data.frame(id = 1:15,
                 grp = rep(1:3,each = 5), 
                 frq = rep(c(3,2,4), each = 5))

In this example, grp refers to the group I want to sample by and frq is the sample size specificied for that group.

Using split , I came up with this possible solution, which gives the desired result but seems rather inefficient :

s <- split(df, df$grp)
lapply(s,function(x) sample_n(x, size = unique(x$frq))) %>% 
      do.call(what = rbind)

Is there a way using just dplyr's group_by and sample_n to do this?

My first thought was:

df %>% group_by(grp) %>% sample_n(size = frq)

but this gives the error:

Error in is_scalar_integerish(size) : object 'frq' not found

Answer 1

This works:

df %>% group_by(grp) %>% sample_n(frq[1])

# A tibble: 9 x 3
# Groups:   grp [3]
     id   grp   frq
  <int> <int> <dbl>
1     3     1     3
2     4     1     3
3     2     1     3
4     6     2     2
5     8     2     2
6    13     3     4
7    14     3     4
8    12     3     4
9    11     3     4

Not sure why it didn't work when you tried it.

Answer 2

library(tidyverse)

# example data 
df <- data.frame(id = 1:15,
                 grp = rep(1:3,each = 5), 
                 frq = rep(c(3,2,4), each = 5))

set.seed(22)

df %>%
  group_by(grp) %>%   # for each group
  nest() %>%          # nest data
  mutate(v = map(data, ~sample_n(data.frame(id=.$id), unique(.$frq)))) %>%  # sample using id values and (unique) frq value
  unnest(v)           # unnest the sampled values

# # A tibble: 9 x 2
#     grp    id
#   <int> <int>
# 1     1     2
# 2     1     5
# 3     1     3
# 4     2     8
# 5     2     9
# 6     3    14
# 7     3    13
# 8     3    15
# 9     3    11

Function sample_n works if you pass as inputs a data frame of ids (not a vector of ids) and one frequency value (for each group).

An alternative version using map2 and generating the inputs for sample_n in advance:

df %>%
  group_by(grp) %>%                                 # for every group
  summarise(d = list(data.frame(id=id)),            # create a data frame of ids
            frq = unique(frq)) %>%                  # get the unique frq value
  mutate(v = map2(d, frq, ~sample_n(.x, .y))) %>%   # sample using data frame of ids and frq value
  unnest(v) %>%                                     # unnest sampled values
  select(-frq)                                      # remove frq column (if needed)

Answer 3

The following answer is not recommended, just shows a different approach without nests/maps that some people might find more comprehensible. Possibly of use to someone working with a smallish data set who wants to do something slightly different to the original question, is a bit scared or doesn't have time to play around with functions they don't really understand, and isn't too worried about efficiency. You just need to recall the behaviour of the original sample function in base R: when provided with a (positive) integer argument x , it outputs a vector randomly permuting the integers from 1:x .

> sample(5)
[1] 5 1 4 2 3

If we had five elements, we could then obtain a random sample of size three by only selecting the positions where 1, 2 and 3 were permuted - in this case we'd pick the second, fourth and fifth elements. All clear? Then similarly we can just do that within each group, assigning random integers from 1 to the group size, and choosing as our sample the places where the random id is less than or equal to the desired sample size for that group.

library(tidyverse)

# The iris data set has three different species
# I want to sample 2, 5 and 3 flowers respectively from each
sample_sizes <- data.frame(
  Species = unique(iris$Species),
  n_to_sample = c(2, 5, 3)
)

iris %>%
  left_join(sample_sizes, by = "Species") %>% # adds column for how many to sample from this species
  group_by(Species) %>% # each species is a group, the size of the group can be found by n()
  mutate(random_id = sample(n())) %>% # give each flower in the group a random id between 1 and n()
  ungroup() %>%
  filter(random_id <= n_to_sample)

Which gave me the output:

# A tibble: 10 x 7
   Sepal.Length Sepal.Width Petal.Length Petal.Width Species    n_to_sample random_id
          <dbl>       <dbl>        <dbl>       <dbl> <fct>            <dbl>     <int>
 1          4.9         3.1          1.5         0.1 setosa               2         1
 2          5.7         4.4          1.5         0.4 setosa               2         2
 3          6.2         2.2          4.5         1.5 versicolor           5         3
 4          6.3         2.5          4.9         1.5 versicolor           5         2
 5          6.4         2.9          4.3         1.3 versicolor           5         5
 6          6           2.9          4.5         1.5 versicolor           5         4
 7          5.5         2.4          3.8         1.1 versicolor           5         1
 8          7.3         2.9          6.3         1.8 virginica            3         1
 9          7.2         3            5.8         1.6 virginica            3         3
10          6.2         3.4          5.4         2.3 virginica            3         2

You can of course pipe through to select(-random_id, -n_to_sample) if you no longer have any use for the final two columns, but I left them in so it's clearer from the output how the code worked.

For the example data given in the question:

library(dplyr)
# example data 
df <- data.frame(id = 1:15,
                 grp = rep(1:3,each = 5), 
                 frq = rep(c(3,2,4), each = 5))

df %>%
  group_by(grp) %>%
  mutate(random_id = sample(n())) %>%
  ungroup() %>%
  filter(random_id <= frq) %>%
  select(-random_id)

# A tibble: 9 x 3
     id   grp   frq
  <int> <int> <dbl>
1     1     1     3
2     2     1     3
3     3     1     3
4     8     2     2
5     9     2     2
6    11     3     4
7    12     3     4
8    13     3     4
9    15     3     4

NB if you're a safety fanatic and x might be zero, and you want to guarantee the length of the output is definitely the same as x , you're better to do sample(seq_len(x)) than sample(x) . That way you get the zero-length vector integer(0) rather than the length-one vector 0 in the case where x is zero. In my code, the mutate will never be working on a row for which n() is zero (if n() were zero then that group is empty so there couldn't be a row there) and this isn't a problem. Just something to be aware of if you're taking this approach somewhere else.

---

Benchmarks for comparison:

f1 <- function(df) { # @AntoniosK with nest and map
  df %>%
    group_by(grp) %>%   # for each group
    nest() %>%          # nest data
    mutate(v = map(data, ~sample_n(data.frame(id=.$id), unique(.$frq)))) %>%  # sample using id values and (unique) frq value
    unnest(v)           # unnest the sampled values
}

f2 <- function(df) { # @AntoniosK with nest and map2
  df %>%
    group_by(grp) %>%                                 # for every group
    summarise(d = list(data.frame(id=id)),            # create a data frame of ids
              frq = unique(frq)) %>%                  # get the unique frq value
    mutate(v = map2(d, frq, ~sample_n(.x, .y))) %>%   # sample using data frame of ids and frq value
    unnest(v) %>%                                     # unnest sampled values
    select(-frq)                                      # remove frq column (if needed)
}

f3 <- function(df) { # @thc
  df %>% group_by(grp) %>% sample_n(frq[1])
}

f4 <- function(df) { # @Silverfish
  df %>%
    group_by(grp) %>%
    mutate(random_id = sample(n())) %>%
    ungroup() %>%
    filter(random_id <= frq) %>%
    select(-random_id)
}


# example data of variable size

df_n <- function(n) {
  data.frame(id = seq_len(3*n),
             grp = rep(1:3,each = n), 
             frq = rep(c(3,2,4), each = n))
}

require(microbenchmark)
microbenchmark(f1(df_n(1e3)), f2(df_n(1e3)), f3(df_n(1e3)), f4(df_n(1e3)),
               f1(df_n(1e6)), f2(df_n(1e6)), f3(df_n(1e6)), f4(df_n(1e6)),
               times=20)

Results strongly favour @thc's df %>% group_by(grp) %>% sample_n(frq[1]) both for data frame with a couple of thousand or couple of million rows. My naive approach takes two or three times as long, and @AntoniosK's faster solution is the one with nest and map2 (worse than mine for smaller data frames but better for the larger ones).

Unit: milliseconds
            expr       min         lq        mean     median         uq       max neval
  f1(df_n(1000))   12.0007   12.27295   12.479760   12.34190   12.46475   13.6403    20
  f2(df_n(1000))    9.5841    9.82185    9.905120    9.87820    9.98865   10.2993    20
  f3(df_n(1000))    1.3729    1.53470    1.593015    1.56755    1.68910    1.8456    20
  f4(df_n(1000))    3.1732    3.21600    3.558855    3.27500    3.57350    5.4715    20
 f1(df_n(1e+06)) 1582.3807 1695.15655 1699.288195 1714.13435 1727.53300 1744.2654    20
 f2(df_n(1e+06))  323.3649  336.94280  407.581130  346.95390  463.69935  911.6647    20
 f3(df_n(1e+06))  216.3265  235.85830  268.756465  247.63620  259.02640  395.9372    20
 f4(df_n(1e+06))  641.5119  663.03510  737.089355  682.69730  803.98205 1132.6586    20