Efficient creation of tridiagonal matrices

Question

How can I create a quadratic band matrix, where I give the diagonal and the first diagonal below and above the diagonal? I am looking for a function like

tridiag(upper, lower, main)

where length(upper)==length(lower)==length(main)-1 and returns, for example,

tridiag(1:3, 2:4, 3:6)

         [,1] [,2] [,3] [,4]
[1,]    3    1    0    0
[2,]    2    4    2    0
[3,]    0    3    5    3
[4,]    0    0    4    6

Is there an efficient way to do it?

Answer 1

This function will do what you want:

tridiag <- function(upper, lower, main){
    out <- matrix(0,length(main),length(main))
    diag(out) <- main
    indx <- seq.int(length(upper))
    out[cbind(indx+1,indx)] <- lower
    out[cbind(indx,indx+1)] <- upper
    return(out)
}

Note that when the index to a matrix is a 2 column matrix, each row in that index is interpreted as the row and column index for a single value in the vector being assigned.