基于data.table按组计算元素抽取概率及优化抽样权重方案咨询

更简洁的概率计算实现 #

你的思路完全没问题，计算每个id组内某行的抽样概率为其cl值在组内的频次占比，这个逻辑非常合理。你的代码可以优化得更简洁，利用data.table的链式操作合并步骤，同时避免不必要的keyby（除非你后续需要按这些列排序）：

library(data.table)
x <- data.table(id=c(1,1,1,2,2,3,4,4,4,4), cl=c("a","b","c","b","b","a","a","b","c","a"))

# 链式操作一步完成概率计算
x[, num := .N, by = .(id, cl)
  ][, den := .N, by = id
    ][, prob := num / den]

如果不需要保留num和den这两个中间列，还可以进一步简化，直接计算prob：

# 不保留中间列的版本
x[, prob := .N / .SD[, .N], by = .(id, cl), .SDcols = NULL]
# 解释：.N 是当前(id,cl)分组的行数（即该cl在id中的出现次数），.SD[, .N] 是当前id分组的总行数

另一种更直观的方式是先计算每个id的总次数，再关联计算概率，适合大型数据集：

den_dt <- x[, .(den = .N), by = id]
x[den_dt, prob := .N / den, on = "id", by = .(id, cl)]

针对抽样目标的优化方案 #

既然你的最终目标是按这个权重在每个id组内抽取一行，其实可以不用显式计算prob列，直接在抽样步骤中动态生成权重，这样能减少内存占用（尤其当数据集很大时）：

set.seed(123) # 设置随机种子确保结果可复现
sampled_data <- x[, {
  # 计算当前id组内每个cl的频次
  cl_freq <- .SD[, .N, by = cl]
  # 为每一行生成对应的权重：cl频次 / 组内总次数
  weights <- cl_freq[.SD, on = "cl", x.N] / .N
  # 按权重抽取一行
  .SD[sample(.N, 1, prob = weights)]
}, by = id]

print(sampled_data)

或者用更简洁的写法，利用table()快速获取组内cl频次：

set.seed(123)
sampled_data <- x[, .SD[sample(.N, 1, prob = table(cl)[cl] / .N)], by = id]

这个抽样逻辑和你预先计算prob列再抽样完全一致：每个行被抽到的概率等于其cl在组内的频次占比。

额外建议 #

如果你的抽样需求是每个id组内抽取一个唯一的cl值（而不是某一行），可以直接对cl进行加权抽样，效率会更高，尤其当组内同一cl有大量重复行时：

set.seed(123)
sampled_cl <- x[, .(sampled_cl = sample(unique(cl), 1, prob = .N / .SD[, .N], by = cl)), by = id]
# 如果需要获取对应行，再关联回原数据
sampled_data <- x[sampled_cl, on = .(id, cl = sampled_cl)]

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