data.table按行基于概率采样哑变量报错问题求助

问题场景 #

当data.table只有2行时，以下代码能正常运行：

library(data.table)
playdata <- data.table(id = c("a","b"), probabilities = c(0.2, 0.3))
playdata[, sampled_dummy := sample(c(0,1),1, prob = probabilities)]

但当表包含3行及以上时，执行直接报错：

library(data.table)
playdata <- data.table(id = c("a","b","c"), probabilities = c(0.2, 0.3, 0.4))
playdata[, sampled_dummy := sample(c(0,1),1, prob = probabilities)]

报错信息：

Error in sample.int(length(x), size, replace, prob) :
incorrect number of probabilities

尝试按行分组的常规 workaround 也会抛出相同错误：

playdata[, sampled_dummy := sample(c(0,1),1, prob = probabilities), by = seq_len(nrow(playdata))]

核心疑问：明明可以强制逐行执行函数，但为什么sample不遵循data.table的常规语法逻辑？它不该默认逐行处理吗？

问题原因 #

问题出在sample函数本身的参数规则，和data.table的语法无关：

• sample(c(0,1), 1, prob = ...)里，prob参数要求传入和待采样向量长度一致的概率值——这里待采样的是c(0,1)（长度为2），所以prob必须是长度为2的向量，分别对应0和1的采样概率。
• 当表只有2行时，probabilities列刚好是长度为2的向量，碰巧满足sample的参数要求，所以代码能跑，但这不是data.table在逐行处理，而是把整个列作为一个长度为2的向量传给了prob。
• 当行数≥3时，probabilities列长度≥3，和待采样向量长度2不匹配，直接触发sample的参数错误。
• 就算按行分组，每个组的probabilities是长度为1的向量，还是和c(0,1)的长度2不匹配，自然还是报错。

正确实现方式 #

要实现每行基于自身的概率值采样1（概率为该行的probabilities）、0（概率为1-该行的probabilities），需要给sample传入正确长度的概率向量，或者换用更适配的函数：

方法1：逐行构造正确概率向量 #

用by = .I实现逐行处理，同时给sample传入c(1-p, p)这个长度为2的概率向量：

playdata[, sampled_dummy := sample(c(0,1), 1, prob = c(1-probabilities, probabilities)), by = .I]

方法2：更高效的向量化实现 #

如果数据量较大，逐行处理效率低，推荐用rbinom函数——它天生支持向量化输入，完全符合data.table的语法习惯：

playdata[, sampled_dummy := rbinom(.N, size = 1, prob = probabilities)]

这里.N是当前表的行数，size=1表示做1次伯努利试验，prob直接传入整个probabilities列，自动按行对应处理，效率比逐行调用sample高很多。

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