按索引依次选取嵌套列表元素的函数效率优化问询

嘿，我明白你的需求——不是简单把嵌套列表拍平，而是按列依次从每个子列表里取元素：先拿所有子列表的第0个元素，再拿第1个，直到所有元素都被提取出来对吧？我来给你分享几个高效的实现方案，比纯Python循环的写法快多了！

核心思路 #

你的需求本质是按索引分组遍历：把所有子列表中索引相同的元素归为一组，然后依次把这些组里的元素拼起来（跳过那些长度不够的子列表的空缺位置）。

低效实现的常见问题 #

如果你的原函数是类似下面这种纯Python双重循环的写法，那确实会在数据量大的时候变慢：

def naive_func(lol):
    max_len = max(len(sublist) for sublist in lol)
    result = []
    for i in range(max_len):
        for sublist in lol:
            if i < len(sublist):
                result.append(sublist[i])
    return result

这种写法需要先计算所有子列表的最大长度，然后两层循环逐个判断索引是否越界，纯Python循环的开销在处理大规模数据时会很明显。

高效优化方案：用itertools.zip_longest #

Python标准库的itertools模块里的zip_longest是专门处理这种不等长可迭代对象的工具，而且它是用C实现的，速度比纯Python循环快很多。

基础版实现 #

from itertools import zip_longest

def optimized_func(lol):
    """
    Takes a list of lists of varying lengths and creates a single list by taking elements
    one by one from each list respectively.
    E.g.: ll = [[1, 2, 3, 4], [5, 6], [7, 8, 9]] 
    Output: [1,5,7,2,6,8,3,9,4]
    """
    result = []
    # zip_longest会把所有子列表按索引对齐，短列表空缺的位置用None填充
    for group in zip_longest(*lol):
        # 过滤掉填充的None，把有效元素加入结果
        result.extend(item for item in group if item is not None)
    return result

测试验证 #

用你给出的例子测试：

ll = [['a', 'b'], ['c'], ['d', 'e', 'f']]
print(optimized_func(ll))  # 输出: ['a', 'c', 'd', 'b', 'e', 'f']

完全符合你的需求！

处理包含None的子列表 #

如果你的嵌套列表里本身就有None元素，那上面的写法会误过滤掉它们。这时候可以用一个独特的标记值来替代默认的fillvalue：

from itertools import zip_longest

def optimized_func_with_none(lol):
    # 创建一个唯一的标记，不会和列表里的任何元素冲突
    sentinel = object()
    result = []
    for group in zip_longest(*lol, fillvalue=sentinel):
        result.extend(item for item in group if item is not sentinel)
    return result

这样即使子列表里有None，也能被正确保留下来。

为什么这个方案更高效？ #

• zip_longest是C实现的迭代器，内部循环的开销远低于纯Python循环
• 迭代器的方式不需要预先计算最大长度，也不会一次性加载所有数据到内存，处理超大列表时更省内存
• extend搭配生成器表达式的写法，比逐个append的效率更高

内容的提问来源于stack exchange，提问作者P. Prunesquallor