如何使用Numpy内置函数实现数组间隔采样并移动起始索引？

需求说明 #

给定一维数组，按固定步长n_sub提取元素生成多个子数组：

import numpy as np

arr = np.array([1,1,1,2,2,2,3,3,3,4,4,4])
n_sub = 3
# 现有循环实现方式
outs = [arr[i::n_sub] for i in range(n_sub)]
# 输出结果：
# [array([1, 2, 3, 4]), array([1, 2, 3, 4]), array([1, 2, 3, 4])]

对于多维数组，同样需要按此规则提取：

arr = np.array([[1,1], [1,1], [1,1], 
                [2,2], [2,2], [2,2], 
                [3,3], [3,3], [3,3],
                [4,4], [4,4], [4,4]
])
# 现有循环方式的输出：
#[array([[1, 1],
#        [2, 2],
#        [3, 3],
#        [4, 4]]),
# array([[1, 1],
#        [2, 2],
#        [3, 3],
#        [4, 4]]),
# array([[1, 1],
#        [2, 2],
#        [3, 3],
#        [4, 4]])]

解决方案：使用Numpy内置的reshape与transpose实现 #

无需循环，可通过reshape调整数组形状后，结合transpose转置维度，直接得到目标结果：

一维数组场景 #

import numpy as np

arr = np.array([1,1,1,2,2,2,3,3,3,4,4,4])
n_sub = 3

# 重塑为(组数, n_sub)的二维数组后转置
reshaped_arr = arr.reshape(-1, n_sub).T
# 转换为数组列表
outs = list(reshaped_arr)
# 输出结果与原循环方式完全一致

多维数组场景 #

arr = np.array([[1,1], [1,1], [1,1], 
                [2,2], [2,2], [2,2], 
                [3,3], [3,3], [3,3],
                [4,4], [4,4], [4,4]
])
n_sub = 3

# 保留原数组列维度，重塑为(组数, n_sub, 列数)的三维数组，再转置调整维度顺序
reshaped_arr = arr.reshape(-1, n_sub, arr.shape[1]).transpose(1, 0, 2)
# 转换为数组列表
outs = list(reshaped_arr)
# 输出结果与原循环方式完全一致

原理说明 #

• 一维场景：reshape(-1, n_sub)将原数组按每n_sub个元素分组，自动推导组数；转置后将每组的对应位置元素聚合为新子数组，完全匹配arr[i::n_sub]的提取逻辑。
• 多维场景：reshape保留原数组的列维度，转置时调整维度顺序为(n_sub, 组数, 列数)，确保每个子数组的结构与原数组一致。

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