R - последовательность расчетов как в прямом, так и в обратном направлении

Question

У меня есть следующий фрейм данных:

id = c("A","A","A","A","A","A","B","B","B","B","B","B","C","C","C","C","C","C")
month = c(1,2,3,4,5,6,1,2,3,4,5,6,1,2,3,4,5,6)
amount = c(0,0,10,0,0,0,0,10,0,10,0,0,0,0,0,10,10,0)

df <- data.frame(id, month, amount)

Что мне нужно сделать (по идентификатору): Рассчитать (с помощью отрицательного числа) разницу в месяцах между нулевыми и ненулевыми строками "сумма" до тех пор, пока "сумма" не станет равной 0. Когда это произойдет, time = 0. Затем, когда "сумма" превышает ноль в последовательности, вычисление (в виде положительного числа) оглянется назад и вычислит разницу в месяцах между ненулевой и исторической нулевой строкой "сумма".

Решение будет выглядеть так:

solution = c(-2,-1,0,1,2,3,-1,0,1,0,1,2,-3,-2,-1,0,0,1)

Как вы, вероятно, можете сказать, довольно сложно найти эту многоуровневую задачу. В идеале ответом будет использование data.table, так как я имею дело с миллионами строк, но dplyr также удовлетворит мои потребности.

Любая помощь приветствуется.

С.

Answer 1

library(data.table)
setDT(DT)

DT[, g := rleid(id, amount != 0)]
DT[, g_id := g - g[1L], by=id]
DT[, v :=  
  if (g_id == 0L) 
    -(.N:1)
  else if (g_id %% 2 == 0)
    1:.N
  else 
    0L
, by=.(id, g_id)]

all.equal(DT$v, solution) # TRUE

Чтобы увидеть, как это работает:

    id month amount  g g_id  v
 1:  A     1      0  1    0 -2
 2:  A     2      0  1    0 -1
 3:  A     3     10  2    1  0
 4:  A     4      0  3    2  1
 5:  A     5      0  3    2  2
 6:  A     6      0  3    2  3
 7:  B     1      0  4    0 -1
 8:  B     2     10  5    1  0
 9:  B     3      0  6    2  1
10:  B     4     10  7    3  0
11:  B     5      0  8    4  1
12:  B     6      0  8    4  2
13:  C     1      0  9    0 -3
14:  C     2      0  9    0 -2
15:  C     3      0  9    0 -1
16:  C     4     10 10    1  0
17:  C     5     10 10    1  0
18:  C     6      0 11    2  1

Вы можете удалить дополнительные столбцы с DT[, c("g", "g_id") := NULL],

Answer 2

С tidyr а также dplyr

library(dplyr)
library(tidyr)

df_new <- df %>% 
    group_by(id) %>% 
    # identify non-zero instances
    mutate(temp = ifelse(amount != 0, month, NA)) %>% 
    # fill down first
    fill(temp, .direction = "down") %>% 
    # fill up after
    fill(temp, .direction = "up") %>% 
    # calculate difference
    mutate(solution = month - temp) %>% 
    # remove temp
    select(-temp)

Результат

#        id month amount solution
#     <fctr> <dbl>  <dbl>    <dbl>
# 1       A     1      0       -2
# 2       A     2      0       -1
# 3       A     3     10        0
# 4       A     4      0        1
# 5       A     5      0        2
# 6       A     6      0        3
# 7       B     1      0       -1
# 8       B     2     10        0
# 9       B     3      0        1
# 10      B     4     10        0
# 11      B     5      0        1
# 12      B     6      0        2
# 13      C     1      0       -3
# 14      C     2      0       -2
# 15      C     3      0       -1
# 16      C     4     10        0
# 17      C     5     10        0
# 18      C     6      0        1