replace diagonal elements in an array

Question

Does anyone know a neat/efficient way to replace diagonal elements in array, similar to the use of diag(x) <- value for a matrix? In other words something like this:

> m<-array(1:27,c(3,3,3))
> for(k in 1:3){
+   diag(m[,,k])<-5
+ }
> m
, , 1

     [,1] [,2] [,3]
[1,]    5    4    7
[2,]    2    5    8
[3,]    3    6    5

, , 2

 [,1] [,2] [,3]
[1,]    5   13   16
[2,]   11    5   17
[3,]   12   15    5

, , 3

     [,1] [,2] [,3]
[1,]    5   22   25
[2,]   20    5   26
[3,]   21   24    5

but without the use of a for loop (my arrays are pretty large and this manipulation will already be within a loop).

Many thanks.

Answer 1

Try this:

with(expand.grid(a = 1:3, b = 1:3), replace(m, cbind(a, a, b), 5))

EDIT:

The question asked for neat/efficient but, of course, those are not the same thing. The one liner here is compact and loop-free but if you are looking for speed I think you will find that the loop in the question is actually the fastest of all the answers.

Answer 2

You can use the following function for that, provided you have only 3 dimensions in your array. You can generalize to more dimensions based on this code, but I'm too lazy to do that for you ;-)

`arraydiag<-` <- function(x,value){
  dims <- dim(x)
  id <- seq_len(dims[1]) +
        dims[2]*(seq_len(dims[2])-1)
  id <- outer(id,(seq_len(dims[3])-1)*prod(dims[1:2]),`+`)
  x[id] <- value
  dim(x) <- dims
  x
}

This works like :

m<-array(1:36,c(3,3,4))
arraydiag(m)<-NA
m

Note that, contrary to the diag() function, this function cannot deal with matrices that are not square. You can look at the source code of diag() to find out how to adapt this code in order it does so.

Answer 3

diagArr <-
function (dim) 
{
    n <- dim[2]
    if(dim[1] != n) stop("expecting first two dimensions to be equal")
    d <- seq(1, n*n, by=n+1)
    as.vector(outer(d, seq(0, by=n*n, length=prod(dim[-1:-2])), "+"))
}

m[diagArr(dim(m))] <- 5

This is written with the intention that it works for dimensions higher than 3 but I haven't tested it in that case. Should be okay though.