Search code examples
rloopsvectorization

Vectorizing an easy loop in R


I am trying to improve my code, but there is a task where I am not able to vectorize a loop. I have 5 sparse matrices, where I have to loop through and add a 1x3-vector resulting from the non-null elements to another matrix. In this example I work with only 1 sparse matrix to keep it simple.

m <- matrix(0, nrow = 3, ncol = 3)
m[1, 2] <- 5
m[3, 3] <- 8
n <- matrix(0, nrow=3, ncol=5)


for (i in 1:nrow(m)){
  for (j in 1:ncol(m)){
    if (!m[i,j]==0){
      n[i, j:(j+2)] <- n[i, j:(j+2)] + rep(m[i, j], 3)
    }
  }
}

The code works, but I have the feeling that there are much better solutions with vectorized functions. I tried and failed with apply functions as the new matrix n has other dimensions as m.

Would be great to get some ideas here.


Solution

  • Here's a vectorized base R solution:

    1. Create the sequence of indices where you want your values to be inserted. Check ?sequence to understand how this works under the hood.
    2. Insert the values
    idx <- sequence(nvec = rep(3, length(m[m != 0])), from = which(m != 0), by = nrow(m))
    n[idx] <- rep(m[m != 0], each = 3)
    n
    
         [,1] [,2] [,3] [,4] [,5]
    [1,]    0    5    5    5    0
    [2,]    0    0    0    0    0
    [3,]    0    0    8    8    8