如何用NumPy高效实现数组列堆叠重组（适配大规模数组）

需求说明 #

现有10行3列的NumPy数组array1，需生成20行2列的数组array2：将array1第一列的每个元素重复2次（对应后续两列的元素数量），后续每一列的元素依次与对应行的第一列元素组成新行，最终堆叠成结果数组。由于实际处理的数组规模极大，需要高效的向量化实现，避免Python循环带来的性能损耗。

原数组array1 #

import numpy as np

array1 = np.array([['Ship.0001.a', -1.52279114723206, 1.52279114723206],
 ['Ship.0002.b', -1.64301979541779, 3.64257907867432],
 ['Ship.0033.g', -35055.33, 35055.33],
 ['Ship.0004.d', -3723.0254, 50663.02],
 ['Cross.0010.a', -53728.453125, 53728.453125],
 ['Cross.0002.e', -4512.5186, 39713.844],
 ['Cross.0033.c', -25680.36, 25680.36],
 ['Cross.0004.d', -0.893001675605774, 0.684107720851898],
 ['Fluid.00010.a', -1.3819944858551, 1.3819944858551],
 ['Fluid.0012.d', -1.53111243247986, 3.5307240486145]])

目标数组array2示例 #

array2 = [['Ship.0001.a', -1.52279114723206],
         ['Ship.0001.a', 1.52279114723206],
         ['Ship.0002.b', -1.64301979541779],
         ['Ship.0002.b', 3.64257907867432],
         ['Ship.0033.g', -35055.33],
         ['Ship.0033.g', 35055.33],
         ['Ship.0004.d', -3723.0254],
         ['Ship.0004.d', 50663.02],
         ['Cross.0010.a', -53728.453125],
         ['Cross.0010.a',  53728.453125],
         ['Cross.0002.e', -4512.5186],
         ['Cross.0002.e', 39713.844],
         ['Cross.0033.c', -25680.36],
         ['Cross.0033.c', 25680.36],
         ['Cross.0004.d', -0.893001675605774],
         ['Cross.0004.d', 0.684107720851898],
         ['Fluid.00010.a', -1.3819944858551],
         ['Fluid.00010.a', 1.3819944858551],
         ['Fluid.0012.d', -1.53111243247986],
         ['Fluid.0012.d', 3.5307240486145]]

高效实现方案 #

利用NumPy的向量化操作（底层C实现，无Python循环）完成转换，性能远高于手动循环，适合大规模数组处理：

# 重复第一列每个元素2次
col1_repeated = np.repeat(array1[:, 0], 2)
# 扁平化后两列数据，得到连续的数值序列
cols2_3_flat = array1[:, 1:].flatten()
# 按列拼接两个数组，生成结果
array2 = np.column_stack((col1_repeated, cols2_3_flat))

方案说明 #

• np.repeat(array1[:,0], 2)：将第一列的每个元素重复2次，直接生成20行的第一列数据，时间复杂度O(n)
• array1[:,1:].flatten()：把后两列的二维数组直接展平为一维，避免循环遍历
• np.column_stack()：将两个一维数组按列拼接成最终的20行2列数组

这种实现完全依赖NumPy的内置优化操作，处理百万级以上数组时，效率比Python循环高100倍以上。

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