如何基于列名向量高效逐行向R函数传递data.table数据？

问题背景 #

现有包含患者ID、心脏病发作年份及多份检查年份的data.table示例如下：

dt <- data.table(patient_ID = c("A","B","C","D"),
                 heart_attack_year = c(2010, 2011, 2016, 2011),
                 exam_1_year = c(2005, 2001, 2002, 2002),
                 exam_2_year = c(2010, 2006, 2010, 2007),
                 exam_3_year = c(2015, 2012, 2011, 2012),
                 exam_4_year = c(2020, 2017, 2017, 2014))

需求为：为每位患者找出离心脏病发作时间最近的检查序号。已实现一个函数接收目标年份、检查标签和检查年份列表，返回最近检查的标签：

which_exam_is_closest <- function(target_year, exam_labels, v){
  return(exam_labels[which.min(abs(target_year - unlist(v)))])
}

原方案通过for循环逐行调用函数实现，但效率低且代码不够简洁，需要更优的、支持自定义列名向量的实现方式。

优化方案 #

方案1：利用data.table的.SD与apply矢量化处理 #

这种方法直接在data.table中逐行计算，代码简洁且效率优于for循环：

# 自定义检查列名向量（可根据不同数据集调整）
exam_cols <- paste0("exam_", 1:4, "_year")
exam_labels <- 1:4

# 生成closest列
dt[, closest := exam_labels[apply(abs(heart_attack_year - .SD), 1, which.min)], .SDcols = exam_cols]

• 说明：.SDcols指定要操作的检查年份列，abs(heart_attack_year - .SD)计算每行中心脏病发作年份与各检查年份的差值绝对值，apply(..., 1, which.min)逐行找出差值最小的位置，最后用exam_labels映射为检查序号。

方案2：转长表后分组计算（适合大数据集） #

将宽表转为长表后，利用data.table的分组操作筛选最近检查，这种方法在数据量较大时效率更高：

# 自定义检查列名向量
exam_cols <- paste0("exam_", 1:4, "_year")

# 转长表：保留患者ID、心脏病发作年份，新增检查标签和检查年份
dt_long <- melt(dt, 
                id.vars = c("patient_ID", "heart_attack_year"),
                measure.vars = exam_cols,
                variable.name = "exam_label",
                value.name = "exam_year")

# 提取检查序号（从列名中提取数字）
dt_long[, exam_label := as.integer(sub("exam_(\\d+)_year", "\\1", exam_label))]

# 分组计算每个患者的最小差值，筛选最近检查
dt_closest <- dt_long[, .SD[which.min(abs(heart_attack_year - exam_year))], by = patient_ID]

# 合并回原表
dt <- dt[dt_closest, on = "patient_ID", closest := exam_label]

• 说明：转长表后，每个患者的多份检查数据变为多行，通过by = patient_ID分组后，用which.min找到每组中差值最小的行，最后将结果合并回原宽表。

方案3：矢量化改造原函数 #

如果希望保留原函数的逻辑，可将其改造为支持矢量化输入的版本，结合data.table的批量处理：

# 矢量化版本的函数
which_exam_is_closest_vec <- function(target_years, exam_labels, exam_year_matrix){
  abs_diff <- abs(target_years - exam_year_matrix)
  return(exam_labels[apply(abs_diff, 1, which.min)])
}

# 生成检查年份矩阵
exam_year_matrix <- as.matrix(dt[, ..exam_cols])

# 生成closest列
dt[, closest := which_exam_is_closest_vec(heart_attack_year, exam_labels, exam_year_matrix)]

• 说明：将检查年份转为矩阵后，函数可一次性处理所有行的差值计算，避免逐行循环的开销。

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