如何通过可迭代对象从NumPy数组中提取指定长度的子数组？

给定以下NumPy数组：

import numpy as np
arr = np.array([0, 1, 2, 3, 4, 5])

已知len(arr)是x的整数倍，需要一个可迭代对象，能从该数组中返回长度为x的子数组。示例代码如下：

x = 2
sub_arrays = [sub_arr for sub_arr in iterable(arr, x)]

预期结果：

sub_arrays = [ np.ndarray( [0, 1] ), np.ndarray( [2, 3] ), np.ndarray( [4, 5] ) ]

注意直接用切片arr[::x]只能取出间隔元素，无法得到子数组：

x = 2
sub_elements = [sub_elem for sub_elem in arr[::x]]

执行结果：

sub_elements = [0, 2, 4]

下面提供几种可行的实现方式：

1. 自定义生成器函数 #

自己写一个生成器作为可迭代对象，逻辑清晰且灵活：

def split_array(arr, chunk_size):
    for i in range(0, len(arr), chunk_size):
        yield arr[i:i+chunk_size]

# 使用示例
arr = np.array([0,1,2,3,4,5])
x = 2
sub_arrays = [sub_arr for sub_arr in split_array(arr, x)]

2. 利用NumPy内置的np.split函数 #

NumPy自带的split函数可以直接将数组分割为指定份数的子数组，因为len(arr)是x的整数倍，分割份数就是len(arr)//x，该函数返回的列表本身支持迭代：

arr = np.array([0,1,2,3,4,5])
x = 2
sub_arrays = [sub_arr for sub_arr in np.split(arr, len(arr)//x)]

3. 先重塑数组再迭代 #

通过reshape将原数组转换为二维数组，其中每一行就是长度为x的子数组，直接迭代二维数组的行即可：

arr = np.array([0,1,2,3,4,5])
x = 2
reshaped_arr = arr.reshape(-1, x)  # -1表示自动计算行数
sub_arrays = [row for row in reshaped_arr]

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