this my first question and I hope to collaborate in the community. I am in a project in which I must fill the NA values ​​with the average of their neighbors from a matrix of ncol = 10 and nrow = 10. I have developed the following code however it is very computationally inefficient:

Code

get_neighbor <- function(matrix, x=1,y=1){

  z <- complex(real = rep(1:nrow(matrix), ncol(matrix)),
               imaginary = rep(1:ncol(matrix), each = nrow(matrix)))
  
  lookup <- lapply(seq_along(z), function(x){
    # calcular la distantancia 
    dist <- which(abs(z - z[x]) < 2)
    # sacar el elemento x del vecindario 
    dist[which(dist != x)]
  })
  index <- (y-1)*(nrow(matrix))+x
  matrix[lookup[[index]]]
  
}

nn_mean <- function(a){
  if(sum(is.na(a))!=ncol(a)*nrow(a)){
    C <- permutations(2, 2, c(1,dim(a)[1]), repeats.allowed = T)
    Borders <- data.frame(matrix(data = 0, ncol = 2, nrow = nrow(a)*2 + ncol(a)*2 - 4))
    Borders[1:nrow(a), 1] <- 1:nrow(a); Borders[1:nrow(a), 2] <- 1
    for(i in 2:(ncol(a)-1)){
      Borders[i + nrow(a) - 1, 2] <- i; Borders[i + 2*(nrow(a) - 1) - 1, 2] <- i
      Borders[i + nrow(a) - 1, 1] <- 1; Borders[i + 2*(nrow(a) - 1) - 1, 1] <- nrow(a)
    }
    Borders[1:ncol(a) + 3*(nrow(a))-4, 2] <- ncol(a)
    Borders[1:ncol(a) + 3*(nrow(a))-4, 1] <- 1:ncol(a)
    
    id <- which(is.na(a), arr.ind = T)
    id <- data.frame(cbind(id, rep(0, nrow(id))))
    
    while(nrow(id)!=0){
      
      for(i in 1:nrow(id)){
        id[i,3] <- sum(is.na(get_neighbor(a, id[i, 1], id[i, 2])))
      }
      
      max_na <- max(id[, 3])
      for(i in 1:(nrow(a)*2 + ncol(a)*2 - 4)){
        if(is.na(a[Borders[i, 1], Borders[i, 2]]) & sum(is.na(get_neighbor(a, Borders[i, 1], Borders[i, 2]))) == 5){
          index <- which(id[,1] == Borders[i, 1] & id[,2] == Borders[i, 2])
          id[index, 3] <- max_na +1
        }
      }
      
      for(i in 1:4){
        if(is.na(a[C[i,1], C[i,2]]) & sum(is.na(get_neighbor(a, C[i, 1], C[i, 2]))) == 3){
          index <- which(id[,1] == C[i, 1] & id[,2] == C[i, 2])
          id[index, 3] <- max_na +1
        }
      }
      
      id <- id[order(id[,3]),]
      index <- which(id[,3]== min(id[,3]))
      for(i in 1:length(index)){
        a[id[i, 1], id[i, 2]] <- mean(get_neighbor(a, id[i, 1], id[i, 2]), na.rm = T)
        if(is.nan(a[id[i, 1], id[i, 2]])){a[id[i, 1], id[i, 2]] <- NA}
      }
      #print(a)
      id <- which(is.na(a), arr.ind = T)
      id <- data.frame(cbind(id, rep(0, nrow(id))))
      
    }
  }
  return(a)
}

example

a <- matrix(data = runif(100, 0, 10), ncol = 10, nrow = 10)
a[a<2] <- NA  

a
          [,1]     [,2]     [,3]     [,4]     [,5]     [,6]     [,7]     [,8]     [,9]    [,10]
 [1,] 2.313512       NA 5.311104 2.832978 9.917106 2.734799 7.309386       NA 4.794476 6.479147
 [2,] 8.855676 7.555101 8.369477 6.346744 7.727896       NA 9.019421 5.061894 9.116066 6.732293
 [3,] 2.948539 7.440258 6.918414 2.155361 3.511407 5.601253       NA 6.561557 9.543535 4.082592
 [4,] 8.455382 9.169974       NA 4.978224 6.202393       NA 9.435753 9.411371       NA 2.128417
 [5,] 7.744456 3.333072 6.975128 5.876849 4.044768 2.948399 5.067653       NA 6.039412 7.350782
 [6,] 8.793417 9.683755 8.053603 7.406450 6.348171 3.122946 9.378282 5.808363 7.923061 6.415419
 [7,] 4.759612 3.431247 4.123641 6.899569 4.464683 6.588431 5.985248 7.962148 6.668238 4.503556
 [8,] 5.992242       NA 7.099657 6.446650       NA 8.448873 5.884961       NA 2.209453 8.103988
 [9,] 6.383036       NA       NA 5.499157 6.972433 3.129470 3.284383 9.150565 8.484186 4.672878
[10,]       NA       NA 4.258936       NA 9.015525       NA       NA       NA       NA 6.639832

nn_mean(a)

          [,1]     [,2]     [,3]     [,4]     [,5]     [,6]     [,7]     [,8]     [,9]    [,10]
 [1,] 2.313512 6.480974 5.311104 2.832978 9.917106 2.734799 7.309386 7.060248 4.794476 6.479147
 [2,] 8.855676 7.555101 8.369477 6.346744 7.727896 6.545895 9.019421 5.061894 9.116066 6.732293
 [3,] 2.948539 7.440258 6.918414 2.155361 3.511407 5.601253 7.111993 6.561557 9.543535 4.082592
 [4,] 8.455382 9.169974 5.855910 4.978224 6.202393 5.258804 9.435753 9.411371 6.587278 2.128417
 [5,] 7.744456 3.333072 6.975128 5.876849 4.044768 2.948399 5.067653 7.580556 6.039412 7.350782
 [6,] 8.793417 9.683755 8.053603 7.406450 6.348171 3.122946 9.378282 5.808363 7.923061 6.415419
 [7,] 4.759612 3.431247 4.123641 6.899569 4.464683 6.588431 5.985248 7.962148 6.668238 4.503556
 [8,] 5.992242 5.298239 7.099657 6.446650 6.056158 8.448873 5.884961 6.203648 2.209453 8.103988
 [9,] 6.383036 5.902524 5.834195 5.499157 6.972433 3.129470 3.284383 9.150565 8.484186 4.672878
[10,] 6.383036 5.731883 4.258936 6.436513 9.015525 5.600453 5.291218 6.689444 7.236865 6.639832

some idea or a function that is efficient?

This can be written in a even short and fast way in R:

nn_impute <- function(dat){
  idx <- which(is.na(dat), TRUE)
  impute <- function(x){
    y <- expand.grid(x[1] + c(-1,0,1), x[2] + c(-1,0,1))
    z <- !(y == 0 | y > nrow(dat) | y> ncol(dat))
    mean(dat[as.matrix(y[rowSums(z) == 2,])], na.rm = TRUE)
  }
  dat[idx] <- apply(idx, 1, impute)
  dat
}

nn_impute(a) ## Returns the filled in values

This code is around 38X faster than the provided code