如何不使用循环对两个多维数组按指定维度应用函数？

我有两个多维数组，想要遍历指定维度并应用一个函数。如果是一维数组的话可以用mapply，但这里需要给两个输入指定操作的维度（margin）。举个简单例子：

创建测试数组：

array_a <- array(runif(60), dim = c(5,4,3))
array_b <- array(runif(60), dim = c(5,4,3))

自定义函数，接收两个二维数组，返回对应列的乘积和：

my_fun <- function(input_1, input_2){
  return(colSums(input_1 * input_2))
}

目前用循环可以实现需求：

output <- array(NA, dim = c(3, 4))
for( i in 1:3){
  output[i,] <- my_fun(array_a[,,i], array_b[,,i])
}

有没有办法不用循环，用向量化操作实现？

方法1：利用purrr::map2处理列表化切片 #

先把两个数组按目标维度拆分成列表，再用map2配对遍历并应用函数，最后合并结果：

library(purrr)

# 将数组按第3维度拆分为列表
list_a <- lapply(1:3, function(i) array_a[,,i])
list_b <- lapply(1:3, function(i) array_b[,,i])

# 遍历配对的列表元素，应用函数后合并为数组
output_map2 <- do.call(rbind, map2(list_a, list_b, my_fun))
dim(output_map2) <- c(3, 4)

方法2：结合aperm与apply重组维度遍历 #

先调整数组维度顺序，把要遍历的维度放到首位，再用apply遍历该维度：

# 调整维度顺序，将原第3维度移到第一个位置
array_a_reshaped <- aperm(array_a, c(3,1,2))
array_b_reshaped <- aperm(array_b, c(3,1,2))

# 遍历重组后的第1维度（原第3维度），应用函数
output_apply <- t(sapply(1:3, function(i) my_fun(array_a_reshaped[i,,], array_b_reshaped[i,,])))

方法3：完全向量化运算（针对当前函数的最优解） #

如果函数逻辑可以直接通过数组原生运算实现，这是效率最高的方式。针对colSums(input_1 * input_2)，可以直接对整个数组做运算：

# 直接计算数组乘积，在第1维度求和后转置调整维度
output_vectorized <- t(colSums(array_a * array_b, dims = 1))

这种方式无需遍历，完全利用R的向量运算特性，适合处理大规模数据。

内容的提问来源于stack exchange，提问作者DataNoob