uni/sose2021/r/w03/P03-2.R

n_diag <- function(mat, n) {
    stopifnot("Matrix not quadratic" = nrow(mat) == ncol(mat))
    stopifnot("Matrix does not have such a (sub)diagonal" = abs(n) < nrow(mat))
    dg <- c()
    for (i in (1:(nrow(mat)-abs(n)))) {
        elem <- if (n > 0) {
            mat[i, i+n]
        } else {
            mat[i+abs(n), i]
        }
        dg <- c(dg, elem)
    }
    return(dg)
}

`n_diag<-` <- function(x, value, n) {
    stopifnot("Matrix not quadratic" = nrow(mat) == ncol(mat))
    stopifnot("Matrix does not have such a (sub)diagonal" = abs(n) < nrow(mat))
    stopifnot("Subdiagonal length and replacement do not match" = length(value) == nrow(mat) - abs(n))
    for (i in (1:(nrow(mat)-abs(n)))) {
        elem <- if (n > 0) {
            mat[i, i+n] <- value[i]
        } else {
            mat[i+abs(n), i] <- value[i]
        }
    }
    return(mat)
}

mat <- matrix(1:16, nrow=4)
print(mat)
print(n_diag(mat, 2))

n_diag(mat, -3) <- 1
print(mat)