How to store for loop outputs of varying sizes in an initialized vector by indices

user_7498839 0
Viewed 34

Problem Statement

Let's say you have the following data:

df <- data.frame(x = rep(0, 10),
                 batch = rep(1:3,c(4,2,4)))

   x batch
1  0     1
2  0     1
3  0     1
4  0     1
5  0     2
6  0     2
7  0     3
8  0     3
9  0     3
10 0     3

You want to loop over the number of unique batches in your dataset and within each batch, apply an algorithm to generate a vector of 1's and 0's. The algorithm is quite long, so for example's sake, let's say it's a random sample:

set.seed(2021)

for(i in seq_len(length(unique(df$batch)))){
  batch_val <- d[which(df$batch == i),]$batch
  #some algorithm to generate 1's and 0's, but using sample() here
  out_x <- sample(c(0,1), length(batch_val), replace = T)
}

You then want to save out_x into the correct indices in df$x. My current rudimentary approach is to explicitly specify indices:

idxb <- 1
idxe <- length(df[which(df$batch == 1),]$batch)

set.seed(2021)
for(i in seq_len(length(unique(df$batch)))){
  batch_val <- d[which(df$batch == i),]$batch
  #some algorithm to generate 1's and 0's, but using sample() here
  out_x <- sample(c(0,1), length(batch_val), replace = T)
  print(out_x)

  #save output
  df$x[idxb:idxe] <- out_x
  
  #update indices
  idxb <- idxb + length(out_X)
  
  if(i < length(unique(df$batch))) {
    idxe <- idxe + length(df[which(df$batch == i+1),]$batch) 
  }
}

Output

The result should look like this:

   x batch
1  0     1
2  1     1
3  1     1
4  0     1
5  1     2
6  1     2
7  1     3
8  0     3
9  1     3
10 1     3

where each iteration of out_x looks like this:

[1] 0 1 1 0
[1] 1 1
[1] 1 0 1 1

Question

What is a faster way to implement this while still using base R?

vector r for-loop
2 Answers

What about using tapply?

out_x <- tapply(df$batch, df$batch, function(x) sample(c(0,1), length(x), replace = T))

#------
$`1`
[1] 0 1 1 1

$`2`
[1] 0 1

$`3`
[1] 1 1 1 1

And then to reassign to df

df$x <- unlist(out_x)

A timing test:

microbenchmark::microbenchmark(f_loop(), f_apply())

#---------
Unit: microseconds
      expr     min       lq     mean  median      uq      max neval
  f_loop() 399.895 425.1975 442.7077 437.754 450.690  612.969   100
 f_apply() 100.449 106.9185 160.5557 110.913 114.909 4867.603   100

Where the functions are defined as

f_loop <- function(){
  
  idxb <- 1
  idxe <- length(df[which(df$batch == 1),]$batch)

  for(i in seq_len(length(unique(df$batch)))){
    
    batch_val <- df[which(df$batch == i),]$batch
    #some algorithm to generate 1's and 0's, but using sample() here
    out_x <- sample(c(0,1), length(batch_val), replace = T)
    #print(out_x)
    
    #save output
    df$x[idxb:idxe] <- out_x
    
    #update indices
    idxb <- idxb + length(out_x)
    
    if(i < length(unique(df$batch))) {
      idxe <- idxe + length(df[which(df$batch == i+1),]$batch) 
    }
  }
  
  return(df$x)
}


f_apply <- function() {
  unlist(tapply(df$batch, df$batch, function(x) sample(c(0,1), length(x), replace = T)))
}
user_124003850

One solution is to remind myself that I can index a vector with a vector!

set.seed(2021)

for(i in seq_len(length(unique(df$batch)))){
  batch_val <- d[which(df$batch == i),]$batch
  #some algorithm to generate 1's and 0's, but using sample() here
  out_x <- sample(c(0,1), length(batch_val), replace = T)
  print(out_x)

  #save output
  idx <- which(df$batch == i)
  df$x[idx] <- out_x
}
user_74988390
Related