如何用data.table/dplyr高效实现层级式两层带放回抽样？

问题背景 #

需要用data.table或tidyverse高效完成两层带放回抽样：

1. 第一层：基于hospital_id，从全部N家医院中带放回抽取N个医院；
2. 第二层：针对每个抽中的医院i，从其下属的M_i个doctor_id中带放回抽取M_i个医生。

当前实现方式存在连接操作缓慢的问题，现有data.table代码如下：

# 获取唯一医院ID数据框
unique_hospitals_df <- unique(hospital_df[, c("hospital_id")])

# 第一层：带放回抽取医院
r_sampled_hospital_ids <- unique_hospitals_df[sample(nrow(unique_hospitals_df),
                      floor(length(unique_hospitals_df) * sample_frac), replace=T), ]

# 关联医生数据
r_df_full <- r_sampled_hospital_ids[,c("hospital_id", "doctor_id")][DT, on = c("hospital_id", "doctor_id"), nomatch = NULL, allow.cartesian = T]

# 第二层：按医院带放回抽取医生
r_DT_resampled <- r_df_full[, .SD[sample(.N, .N, replace=T)], keyby = hospital_id]

示例数据 #

示例数据生成代码：

dt <- data.table(HOSP = rep(LETTERS[1:5], 1:5), 
                 DOC = letters[15:1], 
                 value = 1:15)

生成的数据集：

HOSP DOC value
 1:    A   o     1
 2:    B   n     2
 3:    B   m     3
 4:    C   l     4
 5:    C   k     5
 6:    C   j     6
 7:    D   i     7
 8:    D   h     8
 9:    D   g     9
10:    D   f    10
11:    E   e    11
12:    E   d    12
13:    E   c    13
14:    E   b    14
15:    E   a    15

抽样要求细节：

• 第一层：从HOSP列带放回抽取15次（即原数据集总行数），得到医院序列；
• 第二层：对每个抽中的医院，从其对应的医生组中按医生数量M带放回抽取M次，保留所有列数据。

高效实现方案 #

1. data.table 优化版本 #

核心思路是避免冗余连接操作，直接在原数据上完成分组抽样，再根据第一层抽样结果扩展数据，全程利用data.table的高效分组和索引特性。

写法一：分步生成结果

library(data.table)

# 为原数据标记每个医院的医生数量和行索引
dt[, `:=`(
    hosp_doctor_count = .N,
    row_index = .I
), by = HOSP]

# 第一层：生成带放回的医院抽样序列（示例中抽15次，即nrow(dt)）
sampled_hospitals <- sample(unique(dt$HOSP), size = nrow(dt), replace = TRUE)

# 第二层：对每个抽中的医院，抽取对应行（带放回）并合并
final_result <- rbindlist(lapply(sampled_hospitals, function(hosp) {
    # 获取当前医院的所有行索引
    hosp_rows <- dt[HOSP == hosp, row_index]
    # 带放回抽取M_i次
    dt[sample(hosp_rows, size = dt[HOSP == hosp, hosp_doctor_count], replace = TRUE)]
}))

# 清理临时列
final_result[, c("hosp_doctor_count", "row_index") := NULL]

写法二：简洁向量化版本

library(data.table)

# 先生成每个医院内部的医生抽样模板（带放回抽M_i次）
hospital_doctor_template <- dt[, .SD[sample(.N, .N, replace = TRUE)], by = HOSP]

# 第一层：生成医院抽样序列
sampled_hospitals <- sample(unique(dt$HOSP), size = nrow(dt), replace = TRUE)

# 根据抽样序列扩展模板，得到最终结果
final_result <- hospital_doctor_template[
    HOSP %in% sampled_hospitals
][rep(seq_len(.N), table(sampled_hospitals)[HOSP])]

2. tidyverse 实现方案 #

利用dplyr的分组抽样和行扩展功能，同样避免冗余连接：

library(tidyverse)

# 生成每个医院内部的医生抽样模板（带放回抽M_i次）
hospital_doctor_template <- dt %>%
    group_by(HOSP) %>%
    slice_sample(n = n(), replace = TRUE) %>%
    ungroup()

# 第一层：生成医院抽样序列
sampled_hospitals <- sample(unique(dt$HOSP), size = nrow(dt), replace = TRUE)

# 匹配模板并扩展得到最终结果
final_result <- hospital_doctor_template %>%
    filter(HOSP %in% sampled_hospitals) %>%
    slice(rep(row_number(), table(sampled_hospitals)[HOSP]))

优化核心点 #

• 移除原方案中高成本的笛卡尔连接操作，改用预先生成医院内部抽样模板，再根据第一层抽样结果重复模板的方式，大幅降低内存占用和计算时间；
• 直接在原数据集上完成分组抽样，无需额外生成唯一医院ID数据框；
• 采用向量化操作或高效的列表合并（rbindlist）替代低效循环，提升运行效率。

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