Copying the diagonal of a dataframe/matrix in the first row, by group

Question

Say I have the following matrix mat3, where column 1 is a variable defining 2 groups:

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

In mat3, how do I copy the diagonals of mat1 and mat2 in their respective first rows (i.e. rows 1 and 6)? The pseudocode would be: diag(mat3) by mat3[,1]

I tried the following but it did not work:

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

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

In practice, I need to apply this operation to a dataframe.

Thanks a lot!

Answer 1

do.call(rbind, lapply(split.data.frame(mat3, mat3[,1]), \(x) {
  x[1, ] <- diag(x); x
}))

      [,1] [,2] [,3] [,4] [,5]
 [1,]    1    1    1    1    1
 [2,]    1    1    0    0    0
 [3,]    1    0    1    0    0
 [4,]    1    0    0    1    0
 [5,]    1    0    0    0    1
 [6,]    3    3    3    3    3
 [7,]    3    3    0    0    0
 [8,]    3    0    3    0    0
 [9,]    3    0    0    3    0
[10,]    3    0    0    0    3

If mat3 is a data.frame, you can revise the anonymous function as

\(x) {
  x[1, ] <- diag(as.matrix(x)); x
}