Python中从嵌套结构提取扁平列表的问题

需要从嵌套字典列表中提取扁平的price和ytd列表，现有代码如下：

list_of_data = [{'id':99,
                 'rocketship':{'price':[10, 10, 10, 10, 10], 
                               'ytd':[1, 1, 1.05, 1.1, 1.18]}},
                {'id':898,
                 'rocketship':{'price':[10, 10, 10, 10, 10], 
                               'ytd':[1, 1, 1.05, 1.1, 1.18]}},
                {'id':903,
                 'rocketship':{'price':[20, 20, 20, 10, 10], 
                               'ytd':[1, 1, 1.05, 1.1, 1.18]}},
                {'id':999,
                 'rocketship':{'price':[20, 20, 20, 10, 10], 
                               'ytd':[1, 3, 4.05, 1.1, 1.18]}},
                ]

price, ytd = map(list, zip(*((list_of_data[i]['rocketship']['price'], list_of_data[i]['rocketship']['ytd']) for i in range(0, len(list_of_data)))))

预期输出 #

price = [10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 20, 20, 20, 10, 10, 20, 20, 20, 10, 10]

ytd = [1, 1, 1.05, 1.1, 1.18, 1, 1, 1.05, 1.1, 1.18, 1, 1, 1.05, 1.1, 1.18, 1, 3, 4.05, 1.1, 1.18]

实际问题 #

运行原代码后，得到的price是嵌套列表：

price
Out[19]: 
[[10, 10, 10, 10, 10],
 [10, 10, 10, 10, 10],
 [20, 20, 20, 10, 10],
 [20, 20, 20, 10, 10]]

解决方案 #

原代码的问题在于zip(*...)会将每个字典项的price列表、ytd列表作为整体进行打包，最终得到嵌套结构而非扁平列表。可以通过以下两种方式修正：

方法1：列表推导式（直观易懂） #

price = [p for item in list_of_data for p in item['rocketship']['price']]
ytd = [y for item in list_of_data for y in item['rocketship']['ytd']]

通过双重循环遍历：先遍历list_of_data中的每个字典项，再遍历每个项里的price/ytd列表元素，将元素逐个添加到新列表中，直接得到扁平结构。

方法2：使用itertools.chain（高效简洁） #

from itertools import chain

price = list(chain.from_iterable(item['rocketship']['price'] for item in list_of_data))
ytd = list(chain.from_iterable(item['rocketship']['ytd'] for item in list_of_data))

chain.from_iterable会将多个子列表拼接成一个连续的迭代器，再转换为列表，适合处理数据量较大的场景，效率更高。

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