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

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!

CodePudding user response：

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

CodePudding user response：

Here's an approch that just finds the start of each matrix assuming they are square:

idx <- seq(from=1, to=nrow(mat3), by=ncol(mat3))
for(i in idx) mat3[i, ] <- rep(mat3[i, 1], ncol(mat3))
mat3
#       [,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