如何在含缺失值场景下高效优化R代码，替代循环实现计算

问题场景 #

当矩阵X无缺失值时，向量化矩阵计算的res1效率极高；而sapply()逐行循环的res2结果与res1一致，但计算成本高、效率低：

# 无缺失值时的计算
res1 <- exp(X %*% log(pr.t) + mX %*% log(1 - pr.t)) %*% wts
res2 <- sapply(1:nrow(X), function(i) exp(colSums(X[i,1:5] * log(pr.t[1:5,]), na.rm = TRUE) + colSums(mX[i,1:5] * log(1 - pr.t[1:5,]), na.rm=TRUE))%*% wts)
all.equal(res1[,1], res2)

但实际场景中X常存在缺失值，此时res1会产生NA，res2能得到正确结果，但循环效率无法满足需求：

# 含缺失值时res1失效，res2正确但低效
res1 <- exp(X %*% log(pr.t) + mX %*% log(1 - pr.t)) %*% wts
res2 <- sapply(1:nrow(X), function(i) exp(colSums(X[i,1:5] * log(pr.t[1:5,]), na.rm = TRUE) + colSums(mX[i,1:5] * log(1 - pr.t[1:5,]), na.rm=TRUE))%*% wts)

核心需求：找到一种方法，在X含缺失值时得到与res2一致的结果，同时具备res1的矩阵计算高效性。目前已考虑循环并行处理和Rcpp，希望找到更优实现。

优化方案：向量化缺失值处理 #

通过将缺失值替换为0，将逐行的na.rm=TRUE求和转化为纯矩阵运算，既保证结果正确，又保留向量化计算的高效性：

# 处理缺失值：将NA替换为0，确保缺失项不参与求和
X_clean <- ifelse(is.na(X), 0, X)
mX_clean <- ifelse(is.na(mX), 0, mX)

# 向量化计算，效率与res1一致
res3 <- exp(X_clean %*% log(pr.t) + mX_clean %*% log(1 - pr.t)) %*% wts

# 验证结果一致性
all.equal(res3[,1], res2)

原理说明 #

res2中colSums(..., na.rm=TRUE)的逻辑是忽略缺失值对应的项，将这些项的贡献视为0。通过把X和mX中的NA替换为0，矩阵乘法会自动跳过这些项的求和，与res2的计算逻辑完全一致，同时全程使用向量化操作，效率与res1持平。

测试用示例对象 #

X <- structure(c(0L, 0L, 0L, 0L, 0L, 0L, 0L, 0L, 0L, 0L, 0L, 0L, 0L, 
0L, 0L, 0L, 0L, 0L, 0L, 0L, 0L, 0L, 0L, 1L, 1L), dim = c(5L, 
5L), dimnames = list(NULL, c("Item 1", "Item 2", "Item 3", "Item 4", 
"Item 5")))

pr.t <- structure(c(0.000389840525419771, 0.000389840525419771, 0.000389840525419771, 
0.000389840525419771, 0.000389840525419771, 0.00116782384335194, 
0.00116782384335194, 0.00116782384335194, 0.00116782384335194, 
0.00116782384335194, 0.00293127561410344, 0.00293127561410344, 
0.00293127561410344, 0.00293127561410344, 0.00293127561410344, 
0.00672641421586068, 0.00672641421586068, 0.00672641421586068, 
0.00672641421586068, 0.00672641421586068, 0.0145666908055583, 
0.0145666908055583, 0.0145666908055583, 0.0145666908055583, 0.0145666908055583, 
0.0301824687604691, 0.0301824687604691, 0.0301824687604691, 0.0301824687604691, 
0.0301824687604691, 0.0600531695657659, 0.0600531695657659, 0.0600531695657659, 
0.0600531695657659, 0.0600531695657659, 0.114143103288218, 0.114143103288218, 
0.114143103288218, 0.114143103288218, 0.114143103288218, 0.204278364784018, 
0.204278364784018, 0.204278364784018, 0.204278364784018, 0.204278364784018, 
0.336697623276164, 0.336697623276164, 0.336697623276164, 0.336697623276164, 
0.336697623276164, 0.5, 0.5, 0.5, 0.5, 0.5, 0.663302376723836, 
0.663302376723836, 0.663302376723836, 0.663302376723836, 0.663302376723836, 
0.795721635215982, 0.795721635215982, 0.795721635215982, 0.795721635215982, 
0.795721635215982, 0.885856896711782, 0.885856896711782, 0.885856896711782, 
0.885856896711782, 0.885856896711782, 0.939946830434234, 0.939946830434234, 
0.939946830434234, 0.939946830434234, 0.939946830434234, 0.969817531239531, 
0.969817531239531, 0.969817531239531, 0.969817531239531, 0.969817531239531, 
0.985433309194442, 0.985433309194442, 0.985433309194442, 0.985433309194442, 
0.985433309194442, 0.993273585784139, 0.993273585784139, 0.993273585784139, 
0.993273585784139, 0.993273585784139, 0.997068724385897, 0.997068724385897, 
0.997068724385897, 0.997068724385897, 0.997068724385897, 0.998832176156648, 
0.998832176156648, 0.998832176156648, 0.998832176156648, 0.998832176156648, 
0.99961015947458, 0.99961015947458, 0.99961015947458, 0.99961015947458, 
0.99961015947458), dim = c(5L, 21L))

wts <- c(2.09899121956567e-14, 4.97536860412164e-11, 1.45066128449311e-08, 
1.22535483614825e-06, 4.21923474255167e-05, 0.000708047795481538, 
0.00643969705140876, 0.033952729786543, 0.108392285626419, 0.21533371569506, 
0.270260183572876, 0.21533371569506, 0.10839228562642, 0.0339527297865429, 
0.00643969705140878, 0.000708047795481537, 4.21923474255168e-05, 
1.22535483614826e-06, 1.45066128449309e-08, 4.97536860412161e-11, 
2.09899121956567e-14)

mX <- 1 - X

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