如何提升data.table分组计算PD_3_N及合并步骤的处理速度？

背景与需求 #

处理概率值时涉及两步操作：

1. 将每个id对应的5组概率值分别乘以随机值，用pmin(1, factor*x)确保结果不超过1；
2. 按group_1和group_2分组，计算PD_3_N=1-PROD(1-PD_2_N)。

当前使用lapply实现，但步骤2运行缓慢，合并两步后速度进一步恶化，寻求优化方案。

可复现代码 #

###########
# Dummy data
set.seed(99)
n_col <- 4
size <- 3e6
num_group2 <- 10
vec_1 <- paste0("PD_1_N", (0:n_col))
vec_2 <- paste0("PD_2_N", (0:n_col))
vec_3 <- paste0("PD_3_N", (0:n_col))
id <- rep(seq(1, size, 1), num_group2)
group_1 <- rep(sample(seq(1, size, 1), size=size, replace=TRUE), num_group2)
group_2 <- sort(rep(seq(1, num_group2, 1), size))
factor <- runif(size*num_group2, 0.5, 4)
data <- data.table(id, group_1, group_2, factor)
data[, vec_1] <- data.table(rep(runif(size, 0, 0.5), num_group2), 
                            rep(runif(size, 0, 0.5), num_group2), 
                            rep(runif(size, 0, 0.5), num_group2), 
                            rep(runif(size, 0, 0.5), num_group2), 
                            rep(runif(size, 0, 0.5), num_group2))

当前代码及耗时 #

###############
# lapply step 1
t <- Sys.time()
data[, (vec_2) := lapply(.SD, function(x) pmin(1, factor*x)), .SDcols=vec_1]
Sys.time() - t
# 耗时约1秒

###############
# lapply step 2
t <- Sys.time()
data[, (vec_3) := lapply(.SD, function(x) 1 - prod((1 - x))), 
     by=c("group_1", "group_2"), .SDcols=vec_2]
Sys.time() - t
# 耗时约1.9分钟

######################
# test: 2 steps in one
t <- Sys.time()
data[, (vec_3) := lapply(.SD, function(x) 1 - prod((1 - pmin(1, factor*x)))), 
     by=c("group_1", "group_2"), .SDcols=vec_1]
Sys.time() - t
# 耗时约10分钟

优化方案及解释 #

1. 性能瓶颈分析 #

• 步骤2缓慢原因：lapply逐列调用prod属于标量循环，在分组场景下反复触发小运算，开销累积；未利用data.table的向量化运算优势。
• 合并两步变慢原因：将步骤1的pmin(1, factor*x)嵌入分组循环后，原本全局一次完成的逐行计算，变成每个分组内重复执行，计算量暴增。

2. 优化代码 #

优化步骤2（核心提速）

用对数求和替代直接乘积，既提升速度又避免小数连乘的数值下溢问题：

# 先完成步骤1（保持原高效逻辑）
data[, (vec_2) := lapply(.SD, function(x) pmin(1, factor*x)), .SDcols=vec_1]

# 优化后的步骤2
t <- Sys.time()
data[, (vec_3) := lapply(.SD, function(col) {
  1 - exp(sum(log(1 - col)))
}), by = .(group_1, group_2), .SDcols = vec_2]
Sys.time() - t

进一步优化：矩阵批量运算

将分组内的列转为矩阵，利用apply批量计算，减少循环开销：

data[, (vec_2) := lapply(.SD, function(x) pmin(1, factor*x)), .SDcols=vec_1]

t <- Sys.time()
data[, (vec_3) := {
  mat <- as.matrix(.SD)
  prod_vals <- apply(1 - mat, 2, prod)
  as.list(1 - prod_vals)
}, by = .(group_1, group_2), .SDcols = vec_2]
Sys.time() - t

正确的两步合并方式

若要合并代码，需先全局完成步骤1，再执行分组计算，避免重复运算：

t <- Sys.time()
data[, (vec_2) := lapply(.SD, function(x) pmin(1, factor*x)), .SDcols=vec_1][
  , (vec_3) := lapply(.SD, function(col) 1 - exp(sum(log(1 - col)))), 
  by = .(group_1, group_2), .SDcols = vec_2
]
Sys.time() - t

关键优化点 #

• 用exp(sum(log(1 - col)))替代prod(1 - col)：数值稳定性更强，运算速度更快；
• 避免分组内重复执行全局计算：步骤1是逐行运算，一次性完成效率远高于分组内重复计算；
• 利用矩阵批量运算：减少逐列循环的开销，充分发挥R的向量化运算优势。

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