按LesionResponse分层拆分数据集：同一患者归组及代码问题解决

问题背景 #

我有一个包含127个唯一患者、1052行数据和156个特征的数据集，部分患者对应多行数据。需要将数据集按以下比例拆分：

• 训练集：60%
• 验证集：20%
• 测试集：20%

核心要求：

• 同一患者的所有行必须归为同一组，避免过拟合
• 按二分类结局LesionResponse（0/1）进行分层拆分

原尝试代码出现行丢失、子集间重复行的问题，代码如下：

set.seed(123)

# Get unique patient IDs
patient_ids <- unique(df$PatientID)

# Randomly assign patient IDs to training, validation, or testing group
patient_groups <- rep(c("train", "val", "test"), length.out = length(patient_ids))
patient_groups <- sample(patient_groups)

# Split data by patient group
train_patients <- patient_ids[patient_groups == "train"]
val_patients <- patient_ids[patient_groups == "val"]
test_patients <- patient_ids[patient_groups == "test"]

train_data <- df %>% filter(PatientID %in% train_patients) %>% group_by(PatientID) %>% slice_sample(prop = 0.6, replace = FALSE)
val_data <- df %>% filter(PatientID %in% val_patients) %>% group_by(PatientID) %>% slice_sample(prop = 0.5, replace = FALSE)
test_data <- df %>% filter(PatientID %in% test_patients) %>% group_by(PatientID) %>% slice_sample(prop = 0.5, replace = FALSE)

# Verify that all patients are included in only one group
all_patients <- c(train_patients, val_patients, test_patients)
stopifnot(length(unique(all_patients)) == length(patient_ids))

# Stratify the split based on LesionResponse
prop_train <- sum(train_data$LesionResponse == "1") / nrow(train_data)
prop_val <- sum(val_data$LesionResponse == "1") / nrow(val_data)
prop_test <- sum(test_data$LesionResponse == "1") / nrow(test_data)

# Print proportions of positive LesionResponse in each group
cat("Proportion of positive LesionResponse in training set:", round(prop_train, 2), "\n")
cat("Proportion of positive LesionResponse in validation set:", round(prop_val, 2), "\n")
cat("Proportion of positive LesionResponse in testing set:", round(prop_test, 2), "\n")

原代码问题分析 #

1. 未分层分配患者：随机分组未考虑LesionResponse的分布，容易导致各分组结局比例失衡
2. 不必要的行抽取：slice_sample操作会抽取患者的部分行，直接导致行丢失，违背"同一患者所有行归为一组"的要求
3. 比例控制不严谨：rep+sample的方式无法严格保证60/20/20的患者比例

正确解决方案代码 #

set.seed(123)

# 1. 创建患者-结局映射表（确保同一患者结局一致，若存在冲突需先清理数据）
patient_outcome <- df %>%
  distinct(PatientID, LesionResponse)

# 2. 按结局分层拆分患者：先分60%到训练集
train_patients <- patient_outcome %>%
  group_by(LesionResponse) %>%
  slice_sample(prop = 0.6) %>%
  pull(PatientID)

# 3. 从剩余患者中拆分验证集（20%）和测试集（20%）
remaining_patients <- patient_outcome %>%
  filter(!PatientID %in% train_patients)

val_patients <- remaining_patients %>%
  group_by(LesionResponse) %>%
  slice_sample(prop = 0.5) %>%
  pull(PatientID)

test_patients <- remaining_patients %>%
  filter(!PatientID %in% val_patients) %>%
  pull(PatientID)

# 4. 提取对应分组的所有行数据
train_data <- df %>% filter(PatientID %in% train_patients)
val_data <- df %>% filter(PatientID %in% val_patients)
test_data <- df %>% filter(PatientID %in% test_patients)

# 验证分组正确性
# 检查患者无跨组重复
stopifnot(length(intersect(train_patients, val_patients)) == 0)
stopifnot(length(intersect(train_patients, test_patients)) == 0)
stopifnot(length(intersect(val_patients, test_patients)) == 0)
# 检查所有患者都被分配
stopifnot(length(c(train_patients, val_patients, test_patients)) == length(unique(df$PatientID)))

# 输出各分组的阳性结局比例
prop_train <- mean(train_data$LesionResponse == "1")
prop_val <- mean(val_data$LesionResponse == "1")
prop_test <- mean(test_data$LesionResponse == "1")

cat("训练集阳性结局比例:", round(prop_train, 2), "\n")
cat("验证集阳性结局比例:", round(prop_val, 2), "\n")
cat("测试集阳性结局比例:", round(prop_test, 2), "\n")

代码关键说明 #

• 患者-结局映射：先提取每个患者的唯一结局，确保分层抽样的准确性（若同一患者存在不同结局，需先确认数据合理性）
• 分层抽样：在每个结局组内按比例分配患者，保证各分组的结局比例与原数据一致
• 完整行提取：直接提取患者的所有行，避免行丢失
• 严格分组验证：通过stopifnot确保患者无跨组重复、所有患者都被分配

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