如何为data.table添加至flag列值为1的倒计时列

示例数据 #

dt = structure(list(date = structure(19309:19318, class = c("IDate", "Date")), flag = c(0, 0, 0, 0, 0, 1, 0, 0, 0, 1)), class = c("data.table", "data.frame"), row.names = c(NA, -10L))

方法一：非等值连接（简洁直观） #

利用data.table的非等值连接特性，匹配每行对应的下一个flag=1的行，计算行数差：

# 添加临时行号列
dt[, rn := .I]
# 提取所有flag=1的行号
flag_rows = dt[flag == 1, .(rn_flag = rn)]
# 非等值连接，匹配第一个大于当前行号的flag行，计算剩余行数
dt[flag_rows, on = .(rn < rn_flag), countdown := rn_flag - rn, mult = "first"]
# 将flag=1的行设为0，处理NA值
dt[is.na(countdown) | flag == 1, countdown := 0L]
# 删除临时行号
dt[, rn := NULL]

方法二：从后往前遍历（高效适合大数据） #

通过反向遍历直接计算，避免连接开销，时间复杂度O(n)：

dt[, countdown := 0L]
# 从倒数第二行向前遍历
for (i in (nrow(dt)-1):1) {
  dt[i, countdown := ifelse(flag == 1, 0, dt[i+1, countdown] + 1)]
}

验证结果 #

运行后得到预期输出：

> dt
          date flag countdown
 1: 2022-11-13    0         5
 2: 2022-11-14    0         4
 3: 2022-11-15    0         3
 4: 2022-11-16    0         2
 5: 2022-11-17    0         1
 6: 2022-11-18    1         0
 7: 2022-11-19    0         3
 8: 2022-11-20    0         2
 9: 2022-11-21    0         1
10: 2022-11-22    1         0

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