Efficient create vector encoding summed position-weight of another vector

Question

Suppose I have a vector v and would like to create another vector p of same length that encodes the summed position-weight of each position in v (and 0 for 0's in v). For instance, for v = c(0,1,3,0,1) I get p = c(0, 2, 11, 0, 16). In other words the i'th element in p (that is not 0) becomes v[i]*i plus the previous element in p (that is not 0).

I came up with a method but it seems ugly and I fear its not very time/memory efficient - I need to do this on huge vectors. Any ideas to improve?

fun <- function(v){ 
  res <- NULL
      s = 0
      for(i in 1:length(v)){   
        ifelse(v[i] == 0, res[i] <- 0, {res[i] <- v[i]*i + s; s <- res[i]}) 
      }
      return(res)
}

And then:

> fun(c(0,1,3,0,1))
[1]  0  2 11  0 16

UPDATE:

How to reverse the output of this function, i.e. get back to c(0,1,3,0,1) from c(0,2,11,0,16) ?

Answer 1

One way is to get a cumulative sum of v while multiplying the vector by the position number. The zeroes will remain so and the rest will calculate correctly. Lastly we can multiply the vector by a transformed vector of ones and zeroes to essentially coerce requisite zero values:

cumsum(v*seq_along(v))*+(!!v)
[1]  0  2 11  0 16

For readability, we can also write out:

cumsum(v * seq_along(v)) * as.integer(as.logical(v))

update

Try this for a reversal:

w <- as.logical(r)
r[w] <- c(r[w][1], diff(r[w]))
r / seq_along(r)
[1] 0 1 3 0 1