R语言按组抽取指定数量观测值：代码无报错但结果不符

数据结构 #

> data|>head(20)|>dput()
structure(list(id = c("42190204", "34390202", "34310104", "34310104", 
"34310104", "34310104", "34310104", "34310104", "34310104", "34310104", 
"34310104", "34310104", "34310104", "34310104", "34310104", "34310104", 
"34310104", "34310104", "34310104", "34310104"), sample_size = c(0, 
7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)), row.names = c(NA, 
20L), class = "data.frame")

需求 #

• 已过滤sample_size中的NA值
• 按id分组处理两种场景：
  • 当sample_size小于组内观测数：随机抽取对应数量的观测
  • 当sample_size大于组内观测数：重复现有观测以满足样本量要求
• 最终每个id组的输出行数等于该组的sample_size

原问题代码 #

test<-data|>filter(!is.na(sample_size))|>nest(id)|>
  mutate(
    data_sample = map2(data, sample_size, ~ slice_sample(.x, n = .y))
  ) %>%
  unnest(cols = data_sample)

这段代码无报错，但无法生成符合sample_size要求的分组观测数。

修复方案 #

问题根源 #

1. nest(id)未按id分组嵌套，导致数据聚合逻辑错误
2. slice_sample默认不允许重复抽样，无法处理sample_size大于组内行数的场景
3. 未处理sample_size=0的边界情况

修复后代码 #

library(tidyverse)

test <- data %>%
  filter(!is.na(sample_size)) %>%
  # 确保每个id组的sample_size唯一（若同一id有多个值，取第一个，可按需调整）
  group_by(id) %>%
  mutate(sample_size = first(sample_size)) %>%
  # 按id分组嵌套数据
  nest() %>%
  mutate(
    data_sample = map2(data, sample_size, function(.x, .y) {
      if (.y == 0) {
        # 样本量为0时返回空数据框
        .x[0, ]
      } else {
        # replace=TRUE开启重复抽样，适配样本量大于组内行数的场景
        slice_sample(.x, n = .y, replace = TRUE)
      }
    })
  ) %>%
  unnest(cols = data_sample) %>%
  ungroup()

关键说明 #

• group_by(id) %>% nest()：正确按id分组，每个组对应独立的数据框
• sample_size = first(sample_size)：保证每个id组只有一个样本量参数，避免映射时长度不匹配
• replace = TRUE：允许重复抽取观测，解决样本量大于组内行数的需求
• 增加sample_size=0的判断，返回空行符合业务逻辑

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