如何将二维NumPy数组按指定规则重组为三维数组？

问题描述 #

原二维数组：

arr = [
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
    [20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31],
    [30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41],
    # ... 更多行
    [130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141],
]

期望转换后的三维数组：

[ 
    [ [0, 1], [20, 21], [30, 31], ..., [130, 131]],
    [ [2, 3], [22, 23], [32, 33], ..., [132, 133]],
    [ [4, 5], [24, 25], [34, 35], ..., [134, 135]],
    [ [6, 7], [26, 27], [36, 37], ..., [136, 137]],
    [ [8, 9], [28, 29], [38, 39], ..., [138, 139]],
    [ [10, 11], [30, 31], [40, 41], ..., [140, 141]],
]

之前得到的不符合预期的结果：

[ 
    [0, 1], [2, 3], [4, 5], ..., [10, 11],
    [20, 21], [22, 23], [24, 25], ..., [30, 31],
    # ... 更多行
]

解决方案 #

核心思路是按列分组收集元素：遍历每行中每2个元素的起始索引，跨所有行收集对应位置的切片，最终组成三维数组。

基础循环写法 #

arr = [
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
    [20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31],
    [30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41],
    [130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141],
]

result = []
# 按步长2遍历每行的元素索引，每次取一组2个元素
for i in range(0, len(arr[0]), 2):
    current_group = []
    # 遍历原数组的每一行，收集对应位置的切片
    for row in arr:
        current_group.append(row[i:i+2])
    result.append(current_group)

print(result)

简洁列表推导式写法 #

result = [[row[i:i+2] for row in arr] for i in range(0, len(arr[0]), 2)]

代码说明 #

• 外层循环range(0, len(arr[0]), 2)：控制每组元素的起始位置，步长为2保证每次取2个连续元素。
• 内层循环row[i:i+2]：对每一行截取从索引i到i+2的子数组（左闭右开，正好取2个元素）。
• 最终将所有分组收集起来，得到符合预期的三维数组。

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