纵向研究百万级患者数据集：高效循环与变量衍生优化问询

我从事性传播疾病（如淋病gonorrhoea）及其女性并发症（如异位妊娠ectopic pregnancy）研究。现有两个纵向数据集：

• df_gono：记录各年度组（year_cat）患者的淋病诊断次数（n_gono）
• df_ecto：记录各年度组患者的异位妊娠诊断次数（n_ecto）

需要为df_ecto新增变量gono_status，规则为：患者一旦在某年度组确诊淋病（n_gono>0），该年度组及后续所有年度组取值为1，否则为0。

原循环代码可正常运行，但在包含约100万行、60万唯一患者（pat_id）的真实数据集上运行耗时超3小时，寻求更高效的优化方案。

原示例代码：

df_gono <- data.frame(pat_id = c("A", "A", "A", "A", "B", "B", "B", "B", "C", "C", "C", "C"), 
                      year_cat = c("2000-2005", "2006-2010", "2011-2015", "2016-2020"), 
                      n_gono = c(0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0))

df_ecto <- data.frame(pat_id = c("A", "A", "A", "A", "B", "B", "B", "B", "C", "C", "C", "C"), 
                      year_cat = c("2000-2005", "2006-2010", "2011-2015", "2016-2020"), 
                      n_ecto = c(1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0))

for (i in c(unique(df_ecto$pat_id))) {
    gono <- df_gono[df_gono$pat_id == i, "n_gono"]
    if (any(gono$n_gono > 0)) {
      status_gono <- c(rep(0, each = 1, len = which(gono$n_gono != 0)[1]-1), rep(1, each = 1, len = dim(gono)[1]-(which(gono$n_gono != 0)[1])+1))
    } else {
      status_gono <- c(rep(0, each = 1, len = dim(gono)[1]))
    }
    df_ecto$gono_status[df_ecto$pat_id == i] <- status_gono
}

df_ecto

优化方案1：使用dplyr（tidyverse生态） #

核心思路是先提取每个患者首次确诊淋病的年度组，再通过分组比较年度组顺序生成gono_status，全程采用向量化操作替代循环。

library(dplyr)

# 1. 提取每个患者首次确诊淋病的年度组
first_gono <- df_gono %>%
  filter(n_gono > 0) %>%
  group_by(pat_id) %>%
  slice_min(year_cat) %>% # 按年度组排序取首次确诊记录
  select(pat_id, first_gono_year = year_cat)

# 2. 合并到df_ecto并生成gono_status
df_ecto_optimized <- df_ecto %>%
  left_join(first_gono, by = "pat_id") %>%
  group_by(pat_id) %>%
  mutate(
    # 将年度组转为有序因子，确保顺序判断正确
    year_cat_ordered = factor(year_cat, levels = unique(year_cat), ordered = TRUE),
    first_gono_ordered = factor(first_gono_year, levels = levels(year_cat_ordered), ordered = TRUE),
    gono_status = case_when(
      is.na(first_gono_ordered) ~ 0, # 从未确诊淋病
      year_cat_ordered >= first_gono_ordered ~ 1, # 当前及之后年度组设为1
      TRUE ~ 0
    )
  ) %>%
  ungroup() %>%
  select(-year_cat_ordered, -first_gono_ordered, -first_gono_year) # 移除中间变量

# 查看结果
df_ecto_optimized

优化方案2：使用data.table（大数据集首选） #

data.table针对大数据集做了底层优化，分组、连接操作的速度和内存效率远超base R循环，适合百万级以上数据处理。

library(data.table)

# 转换为data.table格式
setDT(df_gono)
setDT(df_ecto)

# 1. 获取每个患者首次确诊淋病的年度组
first_gono <- df_gono[n_gono > 0, .(first_gono_year = min(year_cat)), by = pat_id]

# 2. 为年度组设置正确排序（确保时间顺序判断准确）
year_levels <- unique(df_gono$year_cat)
df_ecto[, year_cat := factor(year_cat, levels = year_levels, ordered = TRUE)]
first_gono[, first_gono_year := factor(first_gono_year, levels = year_levels, ordered = TRUE)]

# 3. 左连接后生成gono_status，未确诊患者填充0
df_ecto_optimized <- df_ecto[first_gono, on = "pat_id", 
                             gono_status := as.integer(year_cat >= first_gono_year)]
df_ecto_optimized[is.na(gono_status), gono_status := 0]

# 查看结果
df_ecto_optimized

优化原理说明 #

• 取消逐患者循环：原循环需遍历60万次患者，每次做子集提取和赋值，效率极低；优化方案采用向量化分组操作，一次性处理所有患者。
• 利用高效工具包：dplyr和data.table都对分组、连接操作做了C语言级别的底层优化，尤其是data.table在百万级数据集上能将运行时间压缩到分钟级。
• 简化逻辑：通过提取首次确诊时间+比较年度组顺序的方式，替代原循环中生成0/1序列的复杂逻辑，代码更简洁且不易出错。

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