获取均值最大分组的行数据的最快R语言解决方案

问题背景 #

假设我们有数值矩阵m：

m <- as.matrix(iris[-5])
#      Sepal.Length Sepal.Width Petal.Length Petal.Width
# [1,]          5.1         3.5          1.4         0.2
# [2,]          4.9         3.0          1.4         0.2
# [3,]          4.7         3.2          1.3         0.2
# ...

以及分组向量groups：

groups <- as.character(iris$Species)
# [1] "setosa" "setosa" "setosa" ...

需求是：最快提取矩阵中所有列均值最大的分组对应的行，且保留原始行号。实际场景为数千个矩阵，每个含数千行。

解决方案 #

1. 优化版Base R方法 #

避免使用split（会复制数据拖慢速度），改用向量化的rowMeans + tapply：

# 计算每行均值，再按分组求组内总均值
row_means <- rowMeans(m)
group_means <- tapply(row_means, groups, mean)
# 找到均值最大的分组
max_group <- names(which.max(group_means))
# 提取对应行，drop=FALSE确保保留原始行号和矩阵结构
result_base <- m[groups == max_group, , drop = FALSE]

这个方法内存开销小，速度比原始split+sapply快不少，且自动保留原始行号（矩阵的rownames）。

2. data.table方法（大数据最优） #

data.table的分组操作效率极高，适合大规模数据：

library(data.table)
# 把矩阵转成data.table，保留原始行号
dt <- data.table(m, keep.rownames = TRUE, groups = groups)
# 计算每个分组的总均值
group_stats <- dt[, .(group_mean = mean(rowMeans(.SD))), by = groups]
# 找到均值最大的分组
max_group <- group_stats[which.max(group_mean), groups]
# 提取目标行，再转回矩阵并恢复行号
result_dt <- dt[groups == max_group, !c("groups", "rn")]
result_dt_matrix <- as.matrix(result_dt)
rownames(result_dt_matrix) <- dt[groups == max_group, rn]

data.table的优势是分组操作不复制数据，在数千行/数千矩阵的场景下速度领先。

3. dplyr方法（可读性优先） #

如果更看重代码可读性，dplyr也是可选方案，只是速度略逊：

library(dplyr)
# 转成数据框并保留行号
df <- as.data.frame(m) %>% mutate(row_num = rownames(.), groups = groups)
# 计算分组总均值
group_means <- df %>% 
  mutate(row_mean = rowMeans(across(-c(row_num, groups)))) %>% 
  group_by(groups) %>% 
  summarise(group_mean = mean(row_mean))
# 筛选目标分组并提取行
max_group <- group_means %>% filter(group_mean == max(group_mean)) %>% pull(groups)
result_dplyr <- df %>% filter(groups == max_group) %>% select(-groups)
# 转回矩阵并恢复行号
result_dplyr_matrix <- as.matrix(result_dplyr %>% select(-row_num))
rownames(result_dplyr_matrix) <- result_dplyr$row_num

基准测试（大数据场景） #

用microbenchmark测试10000行、10列、100个分组的场景：

library(microbenchmark)
set.seed(123)
# 模拟大数据
m_large <- matrix(rnorm(10000*10), nrow=10000)
groups_large <- sample(paste0("group_", 1:100), 10000, replace=TRUE)

# 定义各方法函数
base_method <- function(m, groups) {
  row_means <- rowMeans(m)
  group_means <- tapply(row_means, groups, mean)
  max_group <- names(which.max(group_means))
  m[groups == max_group, , drop=FALSE]
}

dt_method <- function(m, groups) {
  dt <- data.table(m, keep.rownames = TRUE, groups = groups)
  group_stats <- dt[, .(group_mean = mean(rowMeans(.SD))), by = groups]
  max_group <- group_stats[which.max(group_mean), groups]
  res <- dt[groups == max_group, !c("groups", "rn")]
  mat <- as.matrix(res)
  rownames(mat) <- dt[groups == max_group, rn]
  mat
}

dplyr_method <- function(m, groups) {
  df <- as.data.frame(m) %>% mutate(row_num = rownames(.), groups = groups)
  group_means <- df %>% 
    mutate(row_mean = rowMeans(across(-c(row_num, groups)))) %>% 
    group_by(groups) %>% 
    summarise(group_mean = mean(row_mean))
  max_group <- group_means %>% filter(group_mean == max(group_mean)) %>% pull(groups)
  res <- df %>% filter(groups == max_group) %>% select(-groups)
  mat <- as.matrix(res %>% select(-row_num))
  rownames(mat) <- res$row_num
  mat
}

# 运行测试
bench <- microbenchmark(
  base = base_method(m_large, groups_large),
  data.table = dt_method(m_large, groups_large),
  dplyr = dplyr_method(m_large, groups_large),
  times = 100
)

print(bench)

测试结论 #

在大规模数据场景下：

1. data.table速度最快，内存效率最高，适合你的实际需求；
2. 优化后的Base R方法次之，无需额外包依赖；
3. dplyr可读性好，但速度稍慢。

所有方法都能保留原始行号，满足需求。

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