深度嵌套列表中同形状numpy数组的迭代切片泛化需求问询

嘿，针对你这个处理任意嵌套列表里的numpy数组的需求，我整理了几个实用的方法，帮你轻松迭代或者切片这些底层数组：

递归简直是处理嵌套结构的“天选之子”，逻辑特别清晰：写个函数遍历每个元素，遇到numpy数组就直接返回，遇到嵌套的列表/元组就钻进去继续遍历。

代码示例：

import numpy as np

def flatten_nested_arrays(nested_structure):
    for item in nested_structure:
        if isinstance(item, np.ndarray):
            yield item
        elif isinstance(item, (list, tuple)):  # 根据你的嵌套类型调整，比如还可以加set
            yield from flatten_nested_arrays(item)

用的时候超简单，比如你有这样的嵌套结构：

# 示例：所有底层数组都是(2,2)形状
nested_data = [
    np.array([[1,2],[3,4]]),
    [np.array([[5,6],[7,8]]), [np.array([[9,10],[11,12]]), np.array([[13,14],[15,16]])]],
    (np.array([[17,18],[19,20]]),)
]

• 迭代所有数组：

for arr in flatten_nested_arrays(nested_data):
    print(arr)

• 批量切片（比如取每个数组的第一行）：

first_rows = [arr[0, :] for arr in flatten_nested_arrays(nested_data)]

如果你的嵌套层级特别深（比如超过Python默认的1000层递归限制），递归会触发栈溢出错误。这时候用迭代的方式，用栈来模拟递归过程就稳了：

def flatten_nested_arrays_iterative(nested_structure):
    stack = [nested_structure]
    while stack:
        current = stack.pop()
        if isinstance(current, np.ndarray):
            yield current
        elif isinstance(current, (list, tuple)):
            # 反转一下，保证遍历顺序和递归完全一致（不关心顺序的话可以省略）
            stack.extend(reversed(current))

用法和递归版完全一样，只是底层实现换成了迭代，再也不用担心深度问题~

既然所有底层数组形状都完全相同，不如直接把它们合并成一个更大的numpy数组，这样就能用numpy的矢量化操作，效率比逐个处理高得多！

# 先收集所有数组
all_arrays = list(flatten_nested_arrays(nested_data))
# 合并成形状为 (数组个数, 原行数, 原列数) 的大数组
combined_array = np.stack(all_arrays, axis=0)

# 批量切片示例：取所有数组的第一行
all_first_rows = combined_array[:, 0, :]
# 批量操作示例：给所有数组加1
all_arrays_plus_1 = combined_array + 1

这个方法特别适合需要对所有数组做相同操作的场景，numpy的矢量化运算速度可不是循环能比的~

如果你的嵌套结构里除了列表、元组，还有其他可迭代对象（比如集合），可以稍微修改判断逻辑，兼容更多类型：

def flatten_nested_arrays_general(nested_structure):
    for item in nested_structure:
        if isinstance(item, np.ndarray):
            yield item
        # 排除字符串/字节串（它们也是可迭代的，但我们不想遍历）
        elif hasattr(item, '__iter__') and not isinstance(item, (str, bytes)):
            yield from flatten_nested_arrays_general(item)

不过还是建议明确指定允许的容器类型，避免意外遍历到不想处理的对象~

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