如何向量化R语言中需分组计算的函数？

我有一组包含共享分组标识的多数据点数据集：

group <- rep(c(1:5), times=3)
cost <- rnorm(length(group), 100, 5)
current_score <- rnorm(length(group), 7, 2)
future_score <- current_score*runif(1)

dat <- data.frame(group, cost, current_score, future_score)

同时编写了一个计算整体加权组得分的函数：

wt_score <- function(group, dat)
{
  one_group_dat <- dat[dat$group == group, ]
  wt_score <- sum(one_group_dat$cost * (one_group_dat$current_score - one_group_dat$future_score))/sum(one_group_dat$cost)
  return(wt_score)
}

但实际场景中需处理数万组、数百万数据点，以下循环方式运行过慢：

# 这个太慢了！
dat$wt_score <- 0
for(i in 1:nrow(dat))
{
  dat$wt_score[i] <- wt_score(dat$group[i], dat)
}

请问是否可以向量化上述函数，从而无需使用循环？

当然可以，以下几种高效的分组计算方式，能彻底替代循环，且性能远超逐行循环，适合大数据量场景：

方法1：Base R 原生实现（ave() 函数） #

利用ave()按组批量计算，直接生成对应每行的组加权得分：

# 先计算得分差值
dat$diff <- dat$current_score - dat$future_score
# 按组计算分子和分母的均值，得到加权得分
dat$wt_score <- with(dat, ave(cost * diff, group, FUN = sum) / ave(cost, group, FUN = sum))

方法2：dplyr 分组计算（语法更直观） #

dplyr的分组语法清晰易读，适合复杂逻辑，处理大数据量也有不错的性能：

library(dplyr)

dat <- dat %>%
  mutate(diff = current_score - future_score) %>%
  group_by(group) %>%
  mutate(wt_score = sum(cost * diff) / sum(cost)) %>%
  ungroup()

方法3：data.table （百万级数据最优方案） #

如果数据规模达到百万级甚至更大，data.table的内存效率和运算速度是三者中最优的：

library(data.table)

# 转换为data.table格式
setDT(dat)[, diff := current_score - future_score]
# 按组计算加权得分
dat[, wt_score := sum(cost * diff)/sum(cost), by = group]

性能提升原理 #

原循环的问题在于逐行重复筛选整个数据集，做了大量冗余计算；而上面的方法都是先按组一次性完成所有聚合计算，再将结果映射回对应行，本质是向量化的分组运算，避免了重复筛选的开销，运算效率提升几个数量级。

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