如何編寫一個提取矩陣對角線的函數？

Question

所以我試圖在R中編寫一個函數來提取矩陣的對角線，即像diag（x）那樣操作，但顯然不使用diag（x）。

我不確定從哪裏開始。

Answer 1

適用於非方陣，以及：

diag2 <- function(x){ 
    n <- min(dim(x)) 
    return(x[matrix(rep(1:n, 2), n, 2)]) 
} 

另外，如果你看一下diag你可以看到發生了什麼事情：

if (is.matrix(x)) { 
    if (nargs() > 1L) 
     stop("'nrow' or 'ncol' cannot be specified when 'x' is a matrix") 
    if ((m <- min(dim(x))) == 0L) 
     return(vector(typeof(x), 0L)) 
    y <- x[1 + 0L:(m - 1L) * (dim(x)[1L] + 1)] 
    nms <- dimnames(x) 
    if (is.list(nms) && !any(sapply(nms, is.null)) && identical((nm <- nms[[1L]][seq_len(m)]), 
     nms[[2L]][seq_len(m)])) 
     names(y) <- nm 
    return(y) 
} 
## there's more... 

如果任何人的好奇，我試圖基準，包括@ shayaa的方法...

set.seed(101) 
a <- matrix(runif(1e6), 1e3, 1e3) 

diag3 <- function(x){ 
    x[row(x) == col(x)] 
} 
library(microbenchmark) 
microbenchmark(diag(a), diag2(a), diag3(a)) 
## Unit: microseconds 
##  expr  min  lq  mean  median   uq  max neval 
## diag(a) 23.205 33.915 47.59246 47.3030 58.2355  79.878 100 
## diag2(a) 31.238 37.262 58.03028 57.5665 70.7300 107.546 100 
## diag3(a) 11744.788 12659.595 15425.79847 13874.7265 15054.8285 164130.271 100 

Answer 2

set.seed(1111) 
df <- matrix(rnorm(16),4,4) 
df[row(df)==col(df)] 

也工作，如果您的矩陣不是正方形

set.seed(1111); df <- matrix(rnorm(30),5,6) 
df[row(df)==col(df)]