r - 是否有另一种方法可以在 R 中不使用 for 循环来编写以下函数？

Question

我正在尝试在没有 for 循环的情况下编写以下函数。请注意，我正在尝试复制该功能diag()。

selfdiag <- function(a) {
  j <- c()
  for (i in 1:ncol(a)) {
    j[i] <- a[i, i]
  }
  return(j)
}

考虑一下：

mat <- matrix(rnorm(4), ncol = 2)

该函数selfdiag()应创建与 相同的结果diag()。感谢您对此的任何帮助。

标签： rfunctionfor-loop

Answer 1

您可以使用对角线的行和列索引创建一个数据框，并使用它从矩阵中提取对角线值。

mat <- matrix(rnorm(4), ncol = 2)

diag()这样做的方法 -

diag(mat)
[1] -0.5004046 -0.8785558

另一种方法-

rows_cols <- data.frame(rows = c(1:ncol(mat)), cols = c(1:ncol(mat)))
mat2 <- mat[as.matrix(rows_cols)]
mat2
[1] -0.5004046 -0.8785558

希望这可以帮助！