How to cumsum the elements of a vector under certain condition in R?

Question

My objective is to do a cumulative sum of the elements of a vector and assign the result to each element. But when certain condition is reached, then reset the cumulative sum.

For example:

vector_A <- c(1, 1, -1, -1, -1, 1, -1, -1, 1, -1)

Now, suppose that the condition to reset the cumulative sum is that the next element has a different sign.

Then the desired output is:

vector_B <- c(1, 2, -1, -2, -3, 1, -1, -2, 1, -1)

How can I achieve this?

Answer 1

Using ave:

ave(vector_A, data.table::rleid(sign(A)), FUN = cumsum)
#  [1]  1  2 -1 -2 -3  1 -1 -2  1 -1

---

A formula version of accumulate:

purrr::accumulate(vector_A, ~ ifelse(sign(.x) == sign(.y), .x + .y, .y))
#  [1]  1  2 -1 -2 -3  1 -1 -2  1 -1

Answer 2

You can use a custom function instead of cumsum and accumulate results using e.g. purrr::accumulate:

library(purrr)
vector_A <- c(1, 1, -1, -1, -1, 1, -1, -1, 1, -1)

purrr::accumulate(vector_A, function(a,b) {
  if (sign(a) == sign(b))
    a+b
  else
    b
  })

[1]  1  2 -1 -2 -3  1 -1 -2  1 -1

or if you want to avoid any branch:

purrr::accumulate(vector_A, function(a,b) { b + a*(sign(a) == sign(b))})

[1]  1  2 -1 -2 -3  1 -1 -2  1 -1