如何加速R语言嵌套for循环？多数据集匹配场景优化

问题描述 #

现有两个数据集，其中一个规模极大。尝试通过嵌套for循环为数据集a创建treatment列，但运行速度极慢。曾尝试向量化加速，但因需基于两个数据集做多条件匹配，难以落地。

原循环逻辑：遍历a的每一行，再遍历b的行，若当前a行的treatment已赋值则停止遍历b；若匹配country_code且birth_year落在b行的lower_cutoff和upper_cutoff区间内，则根据birth_year与cohort的大小关系，给treatment赋值1或0。

原代码 #

reform_loop <- function(a, b){
  for(i in 1:nrow(a)) {
    
    for(j in 1:nrow(b)){
      if(!is.na(a[i,"treatment"])){break}
      
      a[i,"treatment"] <- case_when(a[i,"country_code"] == b[j, "country_code"] &
                            a[i,"birth_year"] >= b[j,"cohort"] &
                            a[i,"birth_year"]<= b[j,"upper_cutoff"] ~ 1,
                          
                          a[i,"country_code"] == b[j, "country_code"] &
                            a[i,"birth_year"] < b[j,"cohort"]&
                            a[i,"birth_year"]>= b[j,"lower_cutoff"] ~ 0)
      
    }
  }
  return(a)
}

a <- reform_loop(a, b)

示例数据集 #

# 个体层面数据，含出生年份
a <- data.frame (country_code = c(2,2,2,10,10,10,10,8), 
                               birth_year = c(1920,1930,1940,1970,1980,1990, 2000, 1910))
# 国家层面改革信息，含受影响群体
b <- data.frame(country_code = c(2,10,10,11),
                      lower_cutoff = c(1928, 1975, 1907, 1934),
                      upper_cutoff = c(1948, 1995, 1927, 1948),
                      cohort = c(1938, 1985, 1917, 1942))

期望结果 #

treatment <- c(NA, 0, 1, NA, 0, 1, NA, NA)

加速方案 #

方案1：使用dplyr + fuzzyjoin #

通过模糊连接匹配多条件，再按个体分组保留第一个有效匹配结果，贴合原循环逻辑：

library(dplyr)
library(fuzzyjoin)

a_with_treatment <- a %>%
  # 模糊连接：匹配同国家，且出生年份在上下限区间内的行
  fuzzy_left_join(
    b,
    by = c("country_code" = "country_code",
           "birth_year" = "lower_cutoff",
           "birth_year" = "upper_cutoff"),
    match_fun = list(`==`, `>=`, `<=`)
  ) %>%
  # 计算treatment值
  mutate(treatment = case_when(
    birth_year.x >= cohort ~ 1,
    birth_year.x < cohort ~ 0
  )) %>%
  # 按个体分组，取第一个非NA的treatment（对应原循环break逻辑）
  group_by(country_code.x, birth_year.x) %>%
  slice(which(!is.na(treatment))[1]) %>%
  ungroup() %>%
  # 整理列名并右连接回原数据，保证顺序一致
  select(country_code = country_code.x, birth_year = birth_year.x, treatment) %>%
  right_join(a, by = c("country_code", "birth_year")) %>%
  arrange(match(country_code, a$country_code), match(birth_year, a$birth_year))

# 提取结果
a$treatment <- a_with_treatment$treatment

方案2：使用data.table（推荐大规模数据使用） #

data.table的非等连接性能远超嵌套循环，适合百万级以上数据集：

library(data.table)

# 转换为data.table格式
setDT(a)
setDT(b)

# 非等连接匹配，直接计算treatment
a[b, on = .(country_code, birth_year >= lower_cutoff, birth_year <= upper_cutoff), 
  `:=`(treatment = fifelse(birth_year >= cohort, 1, 0)), 
  by = .EACHI]

# 按个体分组，保留第一个非NA的treatment（对应原循环break逻辑）
a[, treatment := na.omit(treatment)[1], by = .(country_code, birth_year)]

内容的提问来源于stack exchange，提问作者a.avci