基于data.table识别合格治疗周期并计算持续天数

我在data.table中存储了治疗起始与结束的布尔型数据，并非所有起始记录都符合条件，因此设置了qualified_start标识。需要按person_id分组，识别以qualified_start == TRUE为起始、到下一个end == TRUE为结束的治疗周期，计算该周期的天数并将数值填充至治疗结束的行中，每个治疗周期可包含任意数量的观测值。

示例数据 #

person_id = c(1, 1, 1, 1, 1, 2, 2, 2, 2, 2)
start <- c(TRUE, TRUE, FALSE, TRUE, TRUE, TRUE, FALSE, FALSE, TRUE, TRUE)
end <- c(TRUE, FALSE, TRUE, TRUE , FALSE, FALSE, FALSE, TRUE, TRUE, TRUE)
qualified_start <- c(TRUE, TRUE, FALSE, TRUE, FALSE, TRUE, FALSE, FALSE, FALSE, TRUE)

dates = as.Date(c('2023-08-01', '2023-08-02', '2023-08-03', '2023-08-04', '2023-08-05', '2023-07-31','2023-08-01', '2023-08-02', '2023-08-03', '2023-08-04'))

dt <- data.table(dates, person_id, start, end, qualified_start)

数据展示：

dates person_id start   end qualified_start
 1: 2023-08-01         1  TRUE  TRUE            TRUE
 2: 2023-08-02         1  TRUE FALSE            TRUE
 3: 2023-08-03         1 FALSE  TRUE           FALSE
 4: 2023-08-04         1  TRUE  TRUE            TRUE
 5: 2023-08-05         1  TRUE FALSE           FALSE
 6: 2023-07-31         2  TRUE FALSE            TRUE
 7: 2023-08-01         2 FALSE FALSE           FALSE
 8: 2023-08-02         2 FALSE  TRUE           FALSE
 9: 2023-08-03         2  TRUE  TRUE           FALSE
10: 2023-08-04         2  TRUE  TRUE            TRUE

期望输出 #

dates person_id start   end qualified_start duration
 1: 2023-08-01         1  TRUE  TRUE            TRUE        0
 2: 2023-08-02         1  TRUE FALSE            TRUE       NA
 3: 2023-08-03         1 FALSE  TRUE           FALSE        1
 4: 2023-08-04         1  TRUE  TRUE            TRUE        0
 5: 2023-08-05         1  TRUE FALSE           FALSE       NA
 6: 2023-07-31         2  TRUE FALSE            TRUE       NA
 7: 2023-08-01         2 FALSE FALSE           FALSE       NA
 8: 2023-08-02         2 FALSE  TRUE           FALSE        2
 9: 2023-08-03         2  TRUE  TRUE           FALSE       NA
10: 2023-08-04         2  TRUE  TRUE            TRUE        0

解决方案 #

利用data.table的分组和窗口函数可实现需求，具体步骤如下：

1. 按person_id分组，标记每个合格起始点到下一个结束点的周期组；
2. 对每个周期组，提取对应起始日期（该组内第一个qualified_start == TRUE的日期）；
3. 仅在end == TRUE的行计算并填充周期天数，其余行设为NA。

实现代码：

library(data.table)

# 按person_id分组，生成周期标识：每次qualified_start为TRUE时开启新周期
dt[, cycle_id := cumsum(qualified_start), by = person_id]
# 提取每个周期的起始日期
dt[, start_date := first(dates[qualified_start]), by = .(person_id, cycle_id)]
# 计算周期天数，仅填充到end为TRUE的行
dt[, duration := ifelse(end, dates - start_date, NA_integer_), by = .(person_id, cycle_id)]
# 清理临时生成的辅助列（可选操作）
dt[, c("cycle_id", "start_date") := NULL]

运行上述代码后即可得到与期望一致的输出结果。

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