如何高效提取三维数组的(*, y_i, z_i)切片并得到N×3矩阵？

针对三维数组a（维度为(3,4,5)），要高效提取所有(*, y_i, z_i)位置的元素并生成N×3矩阵，可以用以下两种完全向量化的方法，彻底避免低效的逐元素处理：

方法一：数组转置+矩阵索引 #

通过调整数组维度顺序，把需要批量提取的维度前置，再直接索引：

# 定义原数组
a = array(1:60, c(3,4,5))
N = 100000
y = sample(1:4, N, replace = TRUE)
z = sample(1:5, N, replace = TRUE)

# 转置数组，将第一维度（*对应的维度）移到最后
a_trans = aperm(a, c(2, 3, 1))
# 提取(y_i, z_i)对应的切片，直接得到N×3矩阵
result = a_trans[cbind(y, z), ]

原理：转置后数组维度变为(4,5,3)，cbind(y,z)会定位到每个(y_i,z_i)对应的3个元素，提取结果自然按行对应每组(*,y_i,z_i)的元素。

方法二：直接计算线性索引 #

利用R数组列优先存储的特性，直接计算目标元素的内存位置，提取后重新塑形：

# 定义原数组
a = array(1:60, c(3,4,5))
N = 100000
y = sample(1:4, N, replace = TRUE)
z = sample(1:5, N, replace = TRUE)

dim_a = dim(a)
# 计算每个(y_i,z_i)对应的3个元素的线性索引
linear_idx = (z - 1)*dim_a[1]*dim_a[2] + (y - 1)*dim_a[1] + rep(1:dim_a[1], each = N)
# 提取元素并塑形为N×3矩阵（按行排列）
result = matrix(a[linear_idx], ncol = dim_a[1], byrow = TRUE)

原理：线性索引直接对应元素在内存中的存储位置，无需构造大尺寸的多维度索引矩阵，在N极大的场景下效率优势更显著。

方法对比 #

• 转置法代码简洁，可读性高，适合大多数常规场景；
• 线性索引法在超大规模N的场景下略快，减少了转置带来的内存操作开销。

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