在R中基于概率为数据集行分配Group的实现问题

问题场景 #

你有两个数据集：

• df1：包含Period、Group、socio_demo、Probabilities字段，存储各分组下的Group分配概率
• df2：包含Period、ID、socio_demo、Weight字段，需要为每行分配对应Group

需求是按Period+socio_demo分组，根据df1的概率为df2每行随机分配Group，但原代码使用tidyr::uncount(Probabilities)时触发错误：

*Long vectors are not yet supported. Requested output must be less than 2147483647. *

错误原因 #

uncount(Probabilities)会根据概率值重复生成行，当数据量较大或概率为小数时，会生成远超R整数向量最大限制（2^31-1）的行数，直接触发长向量报错。

解决方案 #

不需要通过复制行实现概率抽样，直接在分组内为每行按概率抽取Group即可，以下是两种可行方案：

方案1：分组后直接抽样 #

library(dplyr)
library(purrr)

result <- df2 %>%
  group_by(Period, socio_demo) %>%
  mutate(
    # 提取当前分组对应的Group列表和概率列表
    target_groups = list(df1$Group[df1$Period == cur_group()$Period & df1$socio_demo == cur_group()$socio_demo]),
    target_probs = list(df1$Probabilities[df1$Period == cur_group()$Period & df1$socio_demo == cur_group()$socio_demo]),
    # 按概率为每行分配Group
    Group = map2_chr(target_groups, target_probs, ~sample(.x, size = n(), replace = TRUE, prob = .y))
  ) %>%
  ungroup() %>%
  # 保留需要的列并调整顺序
  select(Period, ID, socio_demo, Weight, Group)

方案2：嵌套df1后关联抽样 #

这种方式更清晰，适合复杂分组场景：

library(dplyr)
library(tidyr)

# 先将df1按Period+socio_demo嵌套，存储每个分组的Group和概率
df1_nested <- df1 %>%
  group_by(Period, socio_demo) %>%
  nest(group_prob = c(Group, Probabilities)) %>%
  ungroup()

# 关联df2后逐行抽样
result <- df2 %>%
  left_join(df1_nested, by = c("Period", "socio_demo")) %>%
  rowwise() %>%
  mutate(
    Group = sample(group_prob$Group, size = 1, prob = group_prob$Probabilities)
  ) %>%
  ungroup() %>%
  select(Period, ID, socio_demo, Weight, Group)

方案优势 #

两种方案都避免了生成海量复制行，直接在原数据行上完成概率抽样，完全避开长向量限制，同时保证分配逻辑符合需求。

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