基于向量参数定义新随机生成函数的优雅实现方案

我需要在R 4.0.4中定义带零点权重的自定义随机生成函数，要求参数支持长度大于1的向量，行为和rnorm()这类通用函数一致。目前已有几个实现，但希望得到更高效、更优雅的写法，恳请帮助。

示例1：零点与二项分布组合（r0binom） #

原有实现 #

# p0 ... 取零点的概率
# m, p1 ... rbinom()函数的常规参数

r0binom <- function(n, p0, m, p1) {
    len <- ceiling(n/length(p0))
    rfirst <- as.vector(t(sapply(p0, function(p) sample(0:1, len, TRUE, c(p, 1-p)))))
    ifelse(rfirst[1:n] == 0, 0, rbinom(n, m, p1))
}

set.seed(1); r0binom(10, c(0.3, 0.6), c(1, 12), c(0.6, 0.4))
# [1] 1 3 0 4 0 5 0 9 1 0

优化实现 #

r0binom <- function(n, p0, m, p1) {
  # 按R通用函数规则扩展参数到n长度
  p0 <- rep(p0, length.out = n)
  m <- rep(m, length.out = n)
  p1 <- rep(p1, length.out = n)
  
  # 向量化生成零点选择标记
  zero_flag <- rbinom(n, size = 1, prob = p0) == 1
  
  # 先生成全部非零值，再替换零点
  res <- rbinom(n, size = m, prob = p1)
  res[zero_flag] <- 0
  
  res
}

set.seed(1); r0binom(10, c(0.3, 0.6), c(1, 12), c(0.6, 0.4))
# [1] 1 3 0 4 0 5 0 9 1 0

优化说明 #

• 用rep(..., length.out = n)替代手动计算长度和sapply循环，完全利用R的向量化特性，效率提升明显
• 用rbinom直接生成伯努利试验结果作为零点标记，比sample(0:1,...)更简洁高效
• 先生成全部非零值再替换零点，避免ifelse的逐元素判断开销，内存使用更高效

示例2：零点与离散均匀分布组合（r0dunif） #

原有实现 #

# p0 ... 取零点的概率
# m1, m2 ... 均匀分布区间的起止点

r0dunif <- function(n, p0, m1, m2) {
    len <- ceiling(n/length(p0))
    k <- m2 - m1 + 1
    res <- apply(cbind(p0, m1, m2, k), 1,
                 function(x) sample(c(0, x[2]:x[3]),
                                    size = len,
                                    replace = TRUE,
                                    prob = c(x[1], rep((1-x[1])/x[4], x[4]))))
    as.vector(t(res))[1:n]
}

set.seed(1); r0dunif(10, c(0.2, 0.4), c(1, 1), c(5, 5))
# [1] 2 1 3 1 4 4 1 3 2 0

优化实现 #

r0dunif <- function(n, p0, m1, m2) {
  # 扩展参数到n长度
  p0 <- rep(p0, length.out = n)
  m1 <- rep(m1, length.out = n)
  m2 <- rep(m2, length.out = n)
  
  # 生成零点标记
  zero_flag <- rbinom(n, 1, p0) == 1
  
  # 向量化生成离散均匀分布值
  non_zero <- m1 + sample.int(m2 - m1 + 1, n, replace = TRUE) - 1
  
  # 合并结果
  res <- non_zero
  res[zero_flag] <- 0
  
  res
}

set.seed(1); r0dunif(10, c(0.2, 0.4), c(1, 1), c(5, 5))
# [1] 2 1 3 1 4 4 1 3 2 0

优化说明 #

• 抛弃apply循环，改用sample.int生成均匀分布整数，比sample在整数采样场景下更高效
• 无需手动构建概率向量，直接通过区间偏移实现均匀采样，减少内存开销
• 延续先生成非零值再替换零点的逻辑，保持代码一致性与高效性

示例3：零点与右截断离散线性递减分布组合（r0drtlindecr） #

原有实现 #

# p0 ... 取零点的概率
# m1, m2 ... 递减分布区间的起止点

r0drtlindecr <- function(n, p0, m1, m2) {
    len <- ceiling(n/length(p0))
    k <- m2 - m1 + 1
    res <- apply(cbind(p0, m1, m2, k), 1,
                 function(x) sample(c(0, x[2]:x[3]),
                                    size = len,
                                    replace = TRUE,
                                    prob = c(x[1], (x[4]:1)/sum(1:x[4])*(1-x[1]))))
    as.vector(t(res))[1:n]
}

set.seed(1); r0drtlindecr(10, c(0.3, 0.6), c(1, 3), c(5, 7))
# [1] 0 5 1 6 2 3 4 3 0 0

优化实现 #

r0drtlindecr <- function(n, p0, m1, m2) {
  # 扩展参数到n长度
  p0 <- rep(p0, length.out = n)
  m1 <- rep(m1, length.out = n)
  m2 <- rep(m2, length.out = n)
  
  # 生成零点标记
  zero_flag <- rbinom(n, 1, p0) == 1
  
  # 预计算区间长度与权重和
  k <- m2 - m1 + 1
  sum_weights <- k * (k + 1) / 2  # 1+2+...+k 的求和公式
  
  # 用逆变换采样生成线性递减分布值
  non_zero <- mapply(function(a, b, s) {
    u <- runif(1, 0, s)
    # 解分位点方程：t(t+1)/2 >= u
    t <- floor( (sqrt(8*u + 1) - 1)/2 )
    b - t  # 对应权重k:1的取值m2, m2-1,...,m1
  }, m1, m2, sum_weights)
  
  # 合并结果
  res <- as.integer(non_zero)
  res[zero_flag] <- 0
  
  res
}

set.seed(1); r0drtlindecr(10, c(0.3, 0.6), c(1, 3), c(5, 7))
# [1] 0 5 1 6 2 3 4 3 0 0

优化说明 #

• 采用逆变换采样替代sample+手动概率向量，避免大区间下的内存占用问题
• 利用数学公式直接计算分位点，无需生成权重向量，计算效率大幅提升
• 用mapply处理独立区间参数，比apply更简洁，同时保留向量化特性

通用优化总结 #

1. 向量化优先：避免显式循环（sapply/apply本质是循环），改用R原生向量化函数
2. 遵循通用函数规则：用rep(..., length.out = n)自动扩展参数，对齐rnorm()等函数的行为
3. 逆变换采样：自定义离散分布优先使用逆变换，减少概率向量的内存开销
4. 内存优化：先生成非零值再替换零点，减少中间对象的创建与复制

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