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) 

Что мне нужно сделать (по идентификатору) является: Calculate (путем отрицательного числа) разница в месяцев между нулевыми и ненулевыми «суммами» строк до тех пор, пока «сумма» не будет равна 0. Когда это произойдет, время = 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

С 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 

Answer 2

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].