How to efficiently sample from a vector of numbers

Viewed 37
data <- data.frame(p1 = runif(10, 0, 1), p2 = runif(10, 0, 1), p3 = runif(10, 0, 1))

I want to sample from (1, 2, 3) based on probabilities (p1, p2, p3). Currently, I've written this in a loop

for(i in 1:nrow(data)){
  data$x <- sample(c(1, 2, 3), size = 1, replace = TRUE, prob = c(data$p1[i], data$p2[i], data$p3[i])
}

Is there a more efficient way to do this? This is quite slow as the number of rows in the data set increases.

2 Answers

We could use rowwise here

library(dplyr)
data %>% 
   rowwise %>%
   mutate(x = sample(c(1, 2, 3), size = 1, replace = TRUE, 
       prob = c_across(matches('^p\\d+$')))) %>% 
   ungroup

-output

# A tibble: 10 x 4
      p1    p2     p3     x
   <dbl> <dbl>  <dbl> <dbl>
 1 0.297 0.602 0.852      2
 2 0.169 0.892 0.481      3
 3 0.589 0.777 0.384      2
 4 0.553 0.940 0.741      2
 5 0.910 0.414 0.0253     2
 6 0.756 0.131 0.963      1
 7 0.193 0.710 0.795      3
 8 0.424 0.807 0.926      3
 9 0.429 0.868 0.859      3
10 0.765 0.992 0.912      2

Or a more efficient approach would be with dapply from collapse

library(collapse)
data$x <- dapply(gvr(data, '^p\\d+$'), MARGIN = 1, FUN = function(x) 
           sample(c(1, 2, 3), size = 1, replace = TRUE, prob = x))

You can use mapply in base R -

data <- transform(data, x = mapply(function(x, y, z) 
                  sample(c(1, 2, 3), 1, prob = c(x, y, z)), p1, p2, p3))

#          p1        p2         p3 x
#1  0.4546779 0.1741450 0.03489881 3
#2  0.4398293 0.8059894 0.27687308 2
#3  0.6514023 0.2697283 0.84126810 3
#4  0.8332548 0.7092438 0.03198509 2
#5  0.1218922 0.9857776 0.86859004 3
#6  0.6887676 0.4206811 0.20454715 2
#7  0.5082115 0.5669147 0.68234409 1
#8  0.1382448 0.1015148 0.89351722 3
#9  0.2095876 0.1667515 0.99908780 2
#10 0.5600337 0.9912484 0.09108034 2
Related