R语言：基于多条件的时间序列ID增量计数器实现问题

需求规则 #

• 初始ID设为1
• 先检测Activity列值为D的行
• 若后续出现Location=2且Activity=A的行，则在下一行将ID加1
• 重复上述逻辑直至数据集末尾

示例数据集（含期望ID） #

df <- data.frame(
  Location = c(2, 3, 3, 2, 1, 2, 2, 2, 1, 3, 3, 1, 2, 3, 2, 2, 1, 2, 3, 2, 1),
  Activity = c("A", "B", "C", "D", "D", "B", "A", "A", "B", "A", "C", "D", "A", "B", "B", "D", "A", "D", "D", "A", "C"),
  Expected_ID = c(1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 4)
)

原函数的问题分析 #

原函数核心错误：

1. 循环中直接给当前行赋值last_id后整体下移ID列，导致ID增量时机完全错位
2. 状态变量重置逻辑与规则不匹配，未准确对应“下一行加1”的要求

正确实现方案 #

方案1：修正后的循环实现 #

increment_id_correct <- function(df) {
  # 初始化ID列，默认全为1
  df$ID <- 1
  current_id <- 1
  waiting_for_a <- FALSE
  
  # 从第2行开始遍历（第1行ID固定为1）
  for (i in 2:nrow(df)) {
    # 检查上一行是否触发ID增量条件
    if (waiting_for_a && df$Location[i-1] == 2 && df$Activity[i-1] == "A") {
      current_id <- current_id + 1
      waiting_for_a <- FALSE
    }
    
    # 检查上一行是否是D，开启等待状态
    if (df$Activity[i-1] == "D") {
      waiting_for_a <- TRUE
    }
    
    # 给当前行赋值ID
    df$ID[i] <- current_id
  }
  
  return(df)
}

# 测试函数
result <- increment_id_correct(df)
# 验证结果是否符合预期
all.equal(result$ID, result$Expected_ID) # 返回TRUE

方案2：向量化实现（dplyr） #

针对大数据集，向量化操作效率更高，用dplyr可简洁实现：

library(dplyr)

df <- df %>%
  # 标记D行
  mutate(is_D = Activity == "D") %>%
  # 标记触发ID增量的A行（Location=2且Activity=A，且之前有未触发的D）
  mutate(trigger = lag(cumsum(is_D) - cumsum(Location == 2 & Activity == "A"), default = 0) > 0 & 
           Location == 2 & Activity == "A") %>%
  # 计算ID增量触发点，累加得到最终ID
  mutate(ID = 1 + cumsum(lead(trigger, default = FALSE))) %>%
  # 清理临时列
  select(-is_D, -trigger)

# 验证结果是否符合预期
all.equal(df$ID, df$Expected_ID) # 返回TRUE

内容的提问来源于stack exchange，提问作者reytla