如何编写按子列表首元素提取嵌套列表元素的Python函数？

核心思路 #

遍历输入的嵌套列表中的每一个子列表，按如下规则处理：

• 取子列表的第一个元素作为待提取的元素数量n
• 从子列表索引为1的位置开始，提取连续n个元素追加到结果列表中
  用Python列表切片特性天然处理边界场景：如果n为0、或者子列表索引1之后的元素总数小于n，切片操作不会抛出索引越界异常，会自动返回符合长度要求的子序列。

代码实现 #

def extract_nested_elements(lst):
    result = []
    for sub_list in lst:
        extract_count = sub_list[0]
        # 从索引1开始提取extract_count个元素
        result.extend(sub_list[1: 1 + extract_count])
    return result

测试验证 #

输入示例参数：

lst = [[1,2,3],[3,4,5,6,8,9], [0,3,4], [1,4,5,6], [4,9,8,6,5,2]]
print(extract_nested_elements(lst))

运行输出：

[2, 4, 5, 6, 4, 9, 8, 6, 5]

和预期输出完全一致。

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