python - 如何在复合列表理解中包含条件？

Question

我想展平列表但保留 NaN。当没有 NaN 时，以下基本代码有效:

l = [[1], [2, 3, 4], [5]]
[item for sublist in l for item in sublist]
> [1, 2, 3, 4, 5]

现在，如果我遇到以下情况，它将中断:

import numpy as np
l = [[1], [2, 3, 4], np.nan, [5]]
[item for sublist in l for item in sublist]
> TypeError: 'float' object is not iterable

这是有道理的，但我需要处理这个边界案例，我不确定如何在上面的列表理解中添加特殊案例条件。我如何修改这种理解以涵盖这种情况？

预期的输出是:

[1, 2, 3, 4, np.nan, 5]    

Answer 1

我会检查子列表是否可迭代。如果不是，则将其包装在一个元组中，然后进行展平:

import numpy as np

l = [[1], [2, 3, 4], np.nan, [5]]

[item for sublist in l for item in (sublist if hasattr(sublist, "__iter__") else [sublist])]

>>> [1, 2, 3, 4, nan, 5]

使用 itertools 中的 chain 会使它更干净一些:

from itertools import chain

list(chain(*(sublist if hasattr(sublist, "__iter__") else (sublist, ) for sublist in l)))
>>> [1, 2, 3, 4, nan, 5]

Source for checking if iterable (StackOverflow)