函数内循环与外部循环行为差异及嵌套列表flatten函数解析

问题背景 #

我们需要实现一个flatten函数，将无限嵌套的整数列表转换为线性列表，示例如下：

• flatten([1, [2, 3, [4]], 5]) → [1, 2, 3, 4, 5]
• flatten([1, [2, 3], [[2], 5], 6]) → [1, 2, 3, 2, 5, 6]
• flatten([[[[9]], [1, 2], [[8]]]) → [9, 1, 2, 8]

初始代码的问题 #

初始编写的flatten函数中，for循环无法正常识别列表类型，代码如下：

def flatten(sp: list, b=[]) -> list:
    if len(sp) == 0:
        print(b)
    for i in sp:
        if type(i) == list:
            flatten(i)
        elif type(i) == int:
            b.append(i)
            sp.remove(i)

但在函数外的循环却能正常识别列表：

for i in [1, [2, 3, [4]], 5]:
    if type(i) == list:
        print('YES')

为什么函数内循环和外部循环表现不同？ #

核心问题出在循环过程中修改了遍历的列表sp：

• 当遇到整数元素时，执行sp.remove(i)会直接改变原列表的长度和元素位置。比如遍历[1, [2,3],5]时，先处理1并将其移除，列表变为[[2,3],5]，此时循环的下一个元素会跳过原列表的第二个元素（子列表），直接取新列表的第二个元素5，导致子列表被完全忽略。
• 外部循环没有修改原列表，每个元素都会被正常遍历到，自然能识别出子列表。
• 另外，初始函数用type(i) == list判断类型，灵活性不如isinstance（无法识别继承自list的子类），但这不是循环失效的主要原因。

正确代码的运行逻辑 #

正确实现的flatten函数能正常工作，代码如下：

def flatten(lst, new_lst = []):
    for i in lst:
        if isinstance(i, list):
            flatten(i)
        else:
            new_lst.append(i)
    return new_lst

它能正常运行的原因：

1. 不修改原遍历列表：遍历过程中仅读取每个元素，不执行remove这类修改操作，确保所有元素都能被处理到。
2. 递归处理子列表：遇到子列表时，递归调用flatten(i)，逐层拆解嵌套结构，子列表内的元素最终都会被收集到new_lst中。
3. 默认参数复用：默认参数new_lst = []在函数第一次定义时初始化，后续递归调用都会复用该列表，实现所有元素的统一收集。不过要注意，这种写法如果多次调用函数不传入new_lst，会保留之前的结果，比如连续两次调用flatten([1,2])会返回[1,2,1,2]，优化写法可以将默认参数设为None并在内部初始化：

def flatten(lst, new_lst=None):
    if new_lst is None:
        new_lst = []
    for i in lst:
        if isinstance(i, list):
            flatten(i, new_lst)
        else:
            new_lst.append(i)
    return new_lst

4. 类型判断更鲁棒：isinstance(i, list)相比type(i) == list，能兼容继承自list的自定义子类，适用性更强。

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