Извлечение недиагонального среза большой матрицы

У меня есть большая матрица nxn и хотелось бы снять диагональные фрагменты разного размера. Например:

1 2 3 4 5 6
1 2 3 4 5 6
1 2 3 4 5 6
1 2 3 4 5 6
1 2 3 4 5 6
1 2 3 4 5 6

Мне нужна функция R, которая при задании матрицы и "ширины диагонального среза" вернет матрицу nxn только этих значений. Итак, для матрицы выше и, скажем, 3, я бы получил:

1 x x x x x
1 2 x x x x
1 2 3 x x x
x 2 3 4 x x
x x 3 4 5 x
x x x 4 5 6

В настоящий момент я использую (прощаю) цикл for, который невероятно медленный:

getDiags<-function(ndiags, cormat){
  resmat=matrix(ncol=ncol(cormat),nrow=nrow(cormat))
  dimnames(resmat)<-dimnames(cormat)
  for(j in 1:ndiags){
    resmat[row(resmat) == col(resmat) + j] <- 
      cormat[row(cormat) == col(cormat) + j]
  }
  return(resmat)
}

Я понимаю, что это очень "не-R" способ решить эту проблему. Есть ли лучший способ сделать это, возможно, используя diag или lower.tri?

Ответ 1

size <- 6
mat <- matrix(seq_len(size ^ 2), ncol = size)


low <- 0
high <- 3

delta <- rep(seq_len(ncol(mat)), nrow(mat)) - 
    rep(seq_len(nrow(mat)), each = ncol(mat))
#or Ben Bolker better alternative
delta <- row(mat) - col(mat)
mat[delta < low | delta > high] <- NA
mat

это работает с матрицами 5000 x 5000 на моей машине

Ответ 2

Если вы хотите использовать upper.tri и lower.tri, вы можете написать такие функции, как эти:

cormat <- mapply(rep, 1:6, 6)

u.diags <- function(X, n) {
  X[n:nrow(X),][lower.tri(X[n:nrow(X),])] <- NA
  return(X)
}

или

l.diags <- function(X, n) {
  X[,n:ncol(X)][upper.tri(X[,n:ncol(X)])] <- NA
  return(X)
}

или

n.diags <- function(X, n.u, n.l) {
  X[n.u:nrow(X),][lower.tri(X[n.u:nrow(X),])] <- NA
  X[,n.l:ncol(X)][upper.tri(X[,n.l:ncol(X)])] <- NA
  return(X)
}

l.diags(cormat, 3)
u.diags(cormat, 3)
n.diags(cormat, 3, 1)