data.table 样本，概率存储在列中

我有一个数据表，其中存储在列中的离散分布的概率。

例如，dt <- data.table(p1 = c(0.5, 0.25, 0.1), p2 = c(0.25, 0.5, 0.1), p3 = c(0.25, 0.25, 0.8))

我想创建一个新的随机变量列，使用同一行中的概率进行采样。在 data.table 语法中，我想象它的工作方式如下：

dt[, sample := sample(1:3, 1, prob = c(p1, p2, p3))]

如果有一个类似于“pmin”和“pmax”的“psample”函数，那么这将起作用。我能够使用 apply 来完成这项工作，缺点是对于我的真实数据集，这需要比我想要的更长的时间。有没有办法使用 data.table 来完成这项工作？下面给出应用解决方案。

dt[, sample := apply(dt, 1, function(x) sample(1:3, 1, prob = x[c('p1', 'p2', 'p3')]))]

如果您正在选择1:n你可以用sampl.int哪个更快。还apply正在上matrix是比较快的。将两者放入一个函数中psamp甚至更快。

所以，试试这个（我添加了dt[, 1:3]这样添加列后就不会失败）：

psamp <- function(x) sample.int(n=3, size=1, prob=x)
dt[, sample :=apply(as.matrix(dt[, 1:3]), 1, psamp)]

为了稍微摆脱apply我们可以Vectorize psamp并使用do.call。另外——如@IRTFM他在回答中建议——我们应该利用.SD symbol.

psampv <- Vectorize(function(p1, p2, p3) sample.int(n=3, size=1, replace=TRUE, prob=c(p1, p2, p3)))
dt[, sample := do.call(psampv, .SD), .SDcols=c('p1','p2','p3')]

正如建议的那样，将性能提高一个数量级以上@Henrik in comments https://stackoverflow.com/questions/73001429/data-table-sample-with-probabilities-stored-in-columns/73008578#comment128946317_73001429我们可以使用Rcpp。我稍微修改了代码这个答案 https://stackoverflow.com/a/59357190/6574038并使用新的Rcpp::sample，这会给出相同的结果base::sample与相同的set.seed.

#include <Rcpp.h>
// [[Rcpp::export]]
Rcpp::IntegerVector sample_matrix1(Rcpp::NumericMatrix x, Rcpp::IntegerVector choice_set) {
  int n = x.nrow();
  Rcpp::IntegerVector result(n);
  for (int i = 0; i < n; ++i) {
    Rcpp::NumericVector z(x(i, Rcpp::_));
    result[i] = Rcpp::sample(choice_set, 1, false, z)[0];
  }
  return result;
}

Rcpp::sourceCpp("sample_matrix1.cpp")

dt[, sample := sample_matrix1(as.matrix(.SD), 1:3), .SDcols=c('p1','p2','p3')] 

基准测试，每次重复 100k*100 次：

Unit: milliseconds
          expr        min         lq       mean     median         uq       max neval cld
      psamp_:= 1195.16708 1259.06558 1327.19581 1311.17878 1349.98905 1515.1187   100   b
     psamp_.SD 1225.90467 1257.37766 1318.74885 1289.27571 1335.07736 1522.3423   100   b
     psamp_set 1181.44985 1256.73204 1320.29317 1301.75657 1335.22009 1491.3870   100   b
 psamp_do.call 1181.93117 1251.45863 1316.23306 1285.85710 1337.06674 1476.8023   100   b
          rcpp   60.73652   67.15291   72.76073   70.47052   73.91629  127.8278   100  a