Нужны ускорения для кода многомерного массива

Question

У меня есть 2 многомерных массива - 4-мерный массив и 3-мерный массив - и некоторый код, чтобы найти максимум 4-мерного массива по измерению и создать индекс для выбора из 3-мерного массива на основе этого. На данный момент это довольно медленно, и я хотел бы ускорить процесс.

Представление:

library(microbenchmark)

# Make some arrays to test with
array4d <- array( runif(5*500*50*5 ,-1,0),
                  dim = c(5, 500, 50, 5) )
array3d <- array( runif(5*500*5, 0, 1),
                        dim = c(5, 500, 5))

# The code of interest
microbenchmark( {
    max_idx <- apply(array4d, c(1,2,3), which.max )
    selections <- list()
    for( i in 1:dim(array4d)[3] ){
        selections[[i]] <- apply(array3d, c(1,2), which.max) == max_idx[ , , i]
    }
})

Любые подсказки приветствуются!

(Побочная проблема - я ' m рассматривает возможность замены which.max на nnet::which.is.max для случайного разрыва связей)

Edit: более быстрое решение благодаря @GKi, но я все еще надеюсь на некоторое ускорение:

max_idx <- apply(array4d, c(1,2,3), which.max)
max_idx2 <- apply(array3d, c(1,2), which.max)
selections <- lapply(seq_len(dim(array4d)[3]), function(i) max_idx2 == max_idx[ , , i])

Answer 1

Можно поставить apply(array3d, c(1,2), which.max) за пределами l oop.

microbenchmark( {
  max_idx <- apply(array4d, c(1,2,3), which.max)
  max_idx2 <- apply(array3d, c(1,2), which.max)
  selections <- lapply(seq_len(dim(array4d)[3]), function(i) max_idx2 == max_idx[ , , i])
},
{
  max_idx <- apply(array4d, c(1,2,3), which.max )
  selections <- list()
  for( i in 1:dim(array4d)[3] ){
    selections[[i]] <- apply(array3d, c(1,2), which.max) == max_idx[ , , i]
  }
})
#      min       lq     mean   median       uq      max neval cld
# 204.1650 228.0010 260.3101 256.0132 271.6664 433.8932   100  a 
# 396.5284 448.4167 495.3885 487.7741 530.9028 693.5601   100   b