R中实现类SAS Retain功能的高效变量创建方法

我有一个约10万行的R data.frame，以下述示例数据说明问题：ID代表患者，date为疾病检测日期，A、B分别为两类检测的阳性标识，数据集已按ID和date排序，每行至少有一项阳性检测。

示例数据 #

data <- data.frame(
  ID = c(1,1,1,1,1,2,2,2,2,2,3,3,3,3,3),
  date = as.Date(c(
    "2022-01-01", "2022-01-05", "2022-02-17", "2022-05-13", "2022-09-20",
    "2021-02-02", "2021-02-06", "2021-04-18", "2021-04-19", "2022-08-21",
    "2020-01-01", "2020-03-29", "2020-04-03", "2020-04-04", "2022-01-01"
  )),
  A = c(1,1,0,1,0,0,0,1,0,1,0,0,0,0,0),
  B = c(0,0,1,0,1,1,1,0,1,0,1,1,1,1,1)
)

print(data)

输出：

ID       date A B
1   1 2022-01-01 1 0
2   1 2022-01-05 1 0
3   1 2022-02-17 0 1
4   1 2022-05-13 1 0
5   1 2022-09-20 0 1
6   2 2021-02-02 0 1
7   2 2021-02-06 0 1
8   2 2021-04-18 1 0
9   2 2021-04-19 0 1
10  2 2022-08-21 1 0
11  3 2020-01-01 0 1
12  3 2020-03-29 0 1
13  3 2020-04-03 0 1
14  3 2020-04-04 0 1
15  3 2022-01-01 0 1

感染判定规则 #

• 患者的首个检测日期即为初始感染日期，感染次数n_infec设为1；
• 若检测A呈阳性（A==1），且当前日期与上次感染日期间隔≥45天，判定为新感染：更新infec_date为当前日期，n_infec加1；
• 若检测B呈阳性（B==1），且当前日期与上次感染日期间隔≥90天，判定为新感染：执行上述相同更新操作；
• 若不满足上述新感染条件，沿用最近的infec_date和n_infec。

期望输出 #

ID       date A B infec_date n_infec
1   1 2022-01-01 1 0 2022-01-01       1
2   1 2022-01-05 1 0 2022-01-01       1
3   1 2022-02-17 0 1 2022-01-01       1
4   1 2022-05-13 1 0 2022-05-13       2
5   1 2022-09-20 0 1 2022-09-20       3
6   2 2021-02-02 0 1 2021-02-02       1
7   2 2021-02-06 0 1 2021-02-02       1
8   2 2021-04-18 1 0 2021-04-18       2
9   2 2021-04-19 0 1 2021-04-18       2
10  2 2022-08-21 1 0 2022-08-21       3
11  3 2020-01-01 0 1 2020-01-01       1
12  3 2020-03-29 0 1 2020-01-01       1
13  3 2020-04-03 0 1 2020-04-03       2
14  3 2020-04-04 0 1 2020-04-03       2
15  3 2022-01-01 0 1 2022-01-01       3

当前低效实现（for循环） #

以下逐行遍历的for循环处理10万行数据时速度极慢：

for(i in 1:nrow(data)){
  if(i==1){
    data[i,"infec_date"] = data[i,"date"]
    data[i,"n_infec"] = 1
  }else if(data[i,"ID"] != data[i-1,"ID"]){
    data[i,"infec_date"] = data[i,"date"]
    data[i,"n_infec"] = 1
  }else{
    if(data[i,"A"] == 1 & data[i,"date"] >= data[i-1,"infec_date"] + 45){
      data[i,"infec_date"] = data[i,"date"]
      data[i,"n_infec"] = data[i-1,"n_infec"] + 1
    }else if(data[i,"B"] == 1 & data[i,"date"] >= (data[i-1,"infec_date"] + 90)){
      data[i,"infec_date"] = data[i,"date"]
      data[i,"n_infec"] = data[i-1,"n_infec"] + 1
    }else{
      data[i,"infec_date"] = data[i-1,"infec_date"]
      data[i,"n_infec"] = data[i-1,"n_infec"]
    }
  }
}

SAS参考实现 #

无法访问SAS，但同类逻辑在SAS中的实现如下：

data new_data;
    set data;
    by id date;
    length infec_date n_infec 8.;
    format infec_date mmddyy10.;
    retain infec_date n_infec;
    if first.id then do;
        infec_date=date;
        n_infec=1;
        end;
    if A=1 and date>=infec_date+45 then do;
        infec_date=date;
        n_infec=n_infec+1;
        end;
    else if B=1 and date>=infec_date+90 then do;
        infec_date=date;
        n_infec=n_infec+1;
        end;
run;

高效实现方案 #

1. 使用dplyr + purrr #

利用group_by按患者分组，结合purrr::accumulate处理累积状态，模拟SAS的retain逻辑：

library(dplyr)
library(purrr)

data_processed <- data %>%
  group_by(ID) %>%
  mutate(
    # 定义累积函数，处理每一行的状态更新
    state = accumulate(
      .x = 1:n(),
      .f = function(prev_state, idx) {
        current_row = cur_data()[idx,]
        if(idx == 1){
          return(list(infec_date = current_row$date, n_infec = 1))
        }
        # 检查新感染条件
        if(current_row$A == 1 && current_row$date >= prev_state$infec_date + 45){
          return(list(infec_date = current_row$date, n_infec = prev_state$n_infec + 1))
        } else if(current_row$B == 1 && current_row$date >= prev_state$infec_date + 90){
          return(list(infec_date = current_row$date, n_infec = prev_state$n_infec + 1))
        } else {
          return(prev_state)
        }
      },
      .init = list(infec_date = NA_Date_, n_infec = NA_integer_)
    ) %>% tail(-1) # 去掉初始的NA状态
  ) %>%
  # 从state列表中提取infec_date和n_infec
  mutate(
    infec_date = map_dbl(state, ~.x$infec_date) %>% as.Date(origin = "1970-01-01"),
    n_infec = map_int(state, ~.x$n_infec)
  ) %>%
  select(-state) %>%
  ungroup()

print(data_processed)

2. 使用data.table（最适合大数据量） #

data.table的按组操作和循环优化对大数据量处理效率极高，模拟SAS的retain逻辑：

library(data.table)

setDT(data)

data_processed <- data[, {
  # 初始化状态变量
  infec_date <- date[1]
  n_infec <- 1L
  # 预分配结果向量
  res_infec <- rep(infec_date, .N)
  res_n <- rep(n_infec, .N)
  
  for(i in 2:.N){
    current_date <- date[i]
    current_A <- A[i]
    current_B <- B[i]
    
    if(current_A == 1 && current_date >= infec_date + 45){
      infec_date <- current_date
      n_infec <- n_infec + 1L
    } else if(current_B == 1 && current_date >= infec_date + 90){
      infec_date <- current_date
      n_infec <- n_infec + 1L
    }
    # 更新结果
    res_infec[i] <- infec_date
    res_n[i] <- n_infec
  }
  
  .(date = date, A = A, B = B, infec_date = res_infec, n_infec = res_n)
}, by = ID]

print(data_processed)

方案说明 #

• data.table的实现效率最高，因为它在内存中直接操作，且按组循环的开销远小于全局for循环；
• dplyr+purrr的实现更偏向函数式编程风格，代码可读性好，适合熟悉tidyverse的用户；
• 两种方案均避免了全局逐行遍历，针对每组患者单独处理，大幅提升10万行数据的处理速度。

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