求矩阵的对角元素

Question

我编写了一个函数来将矩阵的对角元素存储到向量中。但输出并不像我预期的那样。代码是

diagonal <- function(x) {
for( i in nrow(x)){ 
for(j in ncol(x)){
  if(i == j) { 
    a <- x[i,j]
  } 
}
}
print(a)
}

我正在将矩阵传递给函数。代码有什么问题？

Answer 1

我们可以使用diag函数

diag(m1)
#[1] 1 5 9

或者

m1[col(m1)==row(m1)]
#[1] 1 5 9

如果我们使用for循环，我们是按行和列的顺序循环，即1:nrow(x)/1:ncol(x)不是按nrow(x)/ncol(x)。

diagonal <- function(x) {
  a <- numeric(0)
 for( i in 1:nrow(x)){ 
  for(j in 1:ncol(x)){
     if(i == j) { 
       a <-  c(a, x[i,j])
     } 
    }
  }
 a
 }

diagonal(m1)
#[1] 1 5 9

数据

m1 <- matrix(1:9, ncol=3)