按组将数据框/矩阵的对角线值复制到各组首行

问题背景 #

现有如下矩阵mat3，其第1列用于定义2个分组：

mat1 <- diag(1, 5, 5)
mat1[,1] <- 1
mat2 <- diag(3, 5, 5)
mat2[,1] <- 3
mat3 <- rbind(mat1, mat2)
mat3

需求是按第1列的分组，将每个分组（对应原mat1、mat2）的对角线值复制到该分组在mat3中的对应首行（即第1行和第6行），伪代码可表示为diag(mat3) by mat3[,1]。

尝试了以下方法但未成功：

fnc <- function(x) {
  res <- x
  res[1,] <- diag(x)
  res <<- res
}

by(mat3, as.factor(mat3[,1]), fnc)
res

实际需将该操作应用于数据框，以下是可行解决方案：

方案1：Base R 原生实现 #

核心逻辑是拆分分组、单独处理、合并结果，无需额外依赖包：

# 先把矩阵转成数据框（如果输入本身是数据框可跳过这步）
df3 <- as.data.frame(mat3)

# 按第1列拆分数据框
split_groups <- split(df3, df3[, 1])

# 遍历每个分组，替换首行为该组的对角线值
processed_groups <- lapply(split_groups, function(group) {
  # 提取当前分组的对角线值（需转成矩阵取对角线）
  diag_values <- diag(as.matrix(group))
  # 替换首行
  group[1, ] <- diag_values
  return(group)
})

# 合并所有处理后的分组
result_df <- do.call(rbind, processed_groups)
# 重置行名（可选，让行号连续）
rownames(result_df) <- NULL

# 查看结果
result_df

方案2：tidyverse 简洁实现 #

用dplyr的分组语法更直观，适合日常数据处理流程：

library(dplyr)

# 转成数据框（输入是数据框则跳过）
df3 <- as.data.frame(mat3)

result_df <- df3 %>%
  # 按第1列分组
  group_by(across(1)) %>%
  # 对所有列执行替换：首行替换为当前分组的对角线值，其余行保持不变
  mutate(across(everything(), ~ ifelse(row_number() == 1, diag(cur_data()), .))) %>%
  # 取消分组状态
  ungroup()

# 查看结果
result_df

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