为何我实现的my_rgeom函数比R原生rgeom函数更快？

1. 基于逆变换抽样实现的my_rgeom函数 #

几何分布逆变换抽样的核心原理：若$X \sim \text{Geom}(p)$（几何分布），$Y \sim \text{Exp}(-\log(1-p))$（指数分布），则$Z=\text{ceiling}(Y)$与$X$服从相同分布。基于此实现的自定义抽样函数如下：

my_rgeom <- function(N, p) {
    U <- runif(N)
    Y <- log(U) / log(1 - p)
    
    return(ceiling(Y))
}

2. 基准测试代码与结果 #

使用microbenchmark工具对比自定义函数与R原生rgeom的性能，测试代码如下：

library(microbenchmark)
library(ggplot2)
library(gridExtra)

N <- 10^6
prob <- c(0.1, 0.2, 0.3, 0.4, 0.5) 

results_list <- list()

for (p in prob) {
    benchmark_results <- microbenchmark(
        my_rgeom(N, p),
        rgeom(N, p),
        times = 100  
    )
    
    # Store the results in the list
    results_list[[as.character(p)]] <- benchmark_results
}
# Plot the results for each value of p
plots <- lapply(prob, function(p) {
    autoplot(results_list[[as.character(p)]], title = paste("p =", p))
})

# Display the plots
gridExtra::grid.arrange(grobs = plots)

测试结果（单位：毫秒）：

results_list 
$`0.1`
Unit: milliseconds
           expr      min       lq      mean    median       uq      max neval
 my_rgeom(N, p)  63.2907  69.1584  76.95829  73.62385  78.8459 172.4273   100
    rgeom(N, p) 115.0211 123.7101 131.01782 127.05655 134.2484 191.6843   100

$`0.2`
Unit: milliseconds
           expr      min       lq      mean    median        uq      max neval
 my_rgeom(N, p)  63.4618  76.8853  87.02115  85.91125  94.29075 195.5412   100
    rgeom(N, p) 117.3709 130.1152 142.02507 135.24175 148.68975 274.5087   100

$`0.3`
Unit: milliseconds
           expr      min        lq      mean   median       uq      max neval
 my_rgeom(N, p)  66.2584  77.31255  85.50857  81.4676  88.8424 178.2553   100
    rgeom(N, p) 115.3894 129.59360 139.56611 135.2741 146.1852 187.7755   100

$`0.4`
Unit: milliseconds
           expr      min       lq      mean   median       uq      max neval
 my_rgeom(N, p)  68.1903  75.1489  83.53151  80.1580  86.8057 180.0446   100
    rgeom(N, p) 116.3780 121.7696 132.27215 126.9434 137.8039 213.7012   100

$`0.5`
Unit: milliseconds
           expr      min        lq      mean   median       uq      max neval
 my_rgeom(N, p)  70.0728  73.63675  82.20571  79.3234  84.6281 178.6583   100
    rgeom(N, p) 116.5077 119.13275 126.49830 124.5081 130.1617 187.6423   100

从测试结果可见，my_rgeom的性能显著优于原生rgeom。

3. 原生rgeom未采用逆变换抽样的原因 #

• 数值稳定性优先：当$p$非常接近0时，$\log(1-p)$会趋近于0，此时计算$\log(U)/\log(1-p)$容易出现数值下溢或精度丢失；原生rgeom采用的算法（如改进的拒绝抽样或迭代二项抽样）在极端概率值下能保持结果准确性。
• 分布定义兼容性：R中rgeom实现的是**“首次成功前失败次数”的几何分布（取值为0,1,2,...），而my_rgeom实现的是“首次成功时试验次数”**（取值为1,2,3,...）。逆变换抽样需要额外调整逻辑才能适配原生函数的定义，增加了复杂度。
• 历史兼容性继承：R基础统计函数多继承自早期S语言实现，这些实现优先考虑稳定性和广泛适用性而非极致性能，后续版本为兼容旧代码不会轻易替换核心算法。
• 多场景平衡：逆变换抽样在大样本、中等概率下表现优异，但原生函数需要兼顾小样本、极端概率（p接近0或1）等多种场景，当前实现的鲁棒性更适合作为通用函数。

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