在R data.table中如何高效识别可变长度的行块？

我有一个包含status列的大型data.table，status的取值为0、1或2。每行status为1的记录前可存在零行或多行status为0或2的记录，示例如下：

dt <- data.table(id = c(1, 2, 3, 4, 5, 6, 7, 8, 9, 10), status = c(1, 0, 1, 0, 2, 1, 0, 0, 0, 1))

需求是识别出以下分组：分组1（id=1）、分组2（id=2:3）、分组3（id=4:6）、分组4（id=7:10），需要替代循环的高效实现方式。

方案1：反向填充分组标识 #

利用data.table内置的nafill函数实现高效分组，步骤清晰直观：

library(data.table)

dt <- data.table(id = c(1, 2, 3, 4, 5, 6, 7, 8, 9, 10), status = c(1, 0, 1, 0, 2, 1, 0, 0, 0, 1))

# 给每个status=1的行标记行号，其他行设为NA
dt[, group := ifelse(status == 1, .I, NA)]
# 从后往前填充NA，让每个记录对应后续最近的status=1分组
dt[, group := nafill(group, type = "locf", fromLast = TRUE)]
# 将倒序的分组号转为正序
dt[, group := max(group) - group + 1]

# 查看结果
dt

运行后生成的group列就是所需的分组标识，完全匹配要求的分组范围。

方案2：极简反转累加写法 #

用一行代码即可完成，逻辑紧凑高效：

dt[, group := rev(cumsum(rev(status == 1)))]

逻辑拆解：

1. rev(status == 1)：将status是否为1的逻辑向量反转，让每个分组的结束点（status=1）转换为累加起始点
2. cumsum(...)：对反转后的向量做累加，生成倒序的分组编号
3. rev(...)：再次反转向量，得到正序的分组编号，完美契合需求

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