R语言：计算DataFrame每行到下一个V3=1的未来行数

给定如下R语言的DataFrame：

library(tidyverse)
data <- data.frame(
  V1 = c("a", "a", "a", "a", "a", "a", "a", "a", "a", "a", "a", "a", "a", "a", "a", "a", "a", "a"),
  V2 = c(1995, 1996, 1997, 1998, 1999, 2000, 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010, 2011, 2011),
  V3 = c(0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1)
)

需要创建新列V4，规则如下：

• 在V1分组内，若V3=1，则V4=0；
• 若V3=0，计算该行到下一个V3=1的未来连续年份（V2）数量；
• 若后续无V3=1，则V4默认-1。

用户尝试了以下代码，但仅能处理首次V3变为1的情况，多切换场景下计算错误：

data %>%
  arrange(V1, V2) %>%
  group_by(V1) %>%
  mutate(
    V4 = ifelse(V3 == 1, 0, seq_along(V2) - which.max(lead(V3, default = 0) == 1) - 1)
  )

期望输出：

a   1995    0   -8  
a   1996    0   -7  
a   1997    0   -6  
a   1998    0   -5  
a   1999    0   -4  
a   2000    0   -3  
a   2001    0   -2  
a   2002    0   -1  
a   2003    1   0   
a   2004    0   -1  
a   2005    1   0   
a   2006    0   -5  
a   2007    0   -4  
a   2008    0   -3  
a   2009    0   -2  
a   2010    0   -1  
a   2011    1   0   

可以通过标记每个V3=1的位置，反向填充获取下一个1的行号，再计算距离，具体代码如下：

data %>%
  arrange(V1, V2) %>%
  group_by(V1) %>%
  mutate(
    # 标记所有V3=1的行的行号，其余行设为NA
    pos = ifelse(V3 == 1, row_number(), NA),
    # 从下往上填充NA，让每个0行获取到下一个V3=1的行号
    next_one_pos = fill(pos, .direction = "up") %>% replace_na(0),
    # 根据规则计算V4值
    V4 = case_when(
      V3 == 1 ~ 0,
      next_one_pos == 0 ~ -1,
      TRUE ~ -(next_one_pos - row_number())
    )
  ) %>%
  select(-pos, -next_one_pos) # 移除辅助列

代码解释 #

1. 标记位置：pos列记录所有V3=1行的行号，非1行设为NA；
2. 反向填充：fill(pos, .direction = "up")从下往上将NA替换为最近的下方V3=1的行号，确保每个0行都能拿到下一个1的位置；
3. 计算V4：
   • 若V3=1，直接赋值0；
   • 若next_one_pos为0，说明该行之后没有V3=1的行，赋值-1；
   • 其余情况计算当前行号与下一个1的行号之差，取负数得到未来的年份数量（例如当前行是第8行，下一个1在第9行，差为1，取负后为-1，匹配期望输出）。

运行代码后即可得到符合要求的结果。

内容的提问来源于stack exchange，提问作者flâneur