I would like to flatten a list but keep NaN
s. The following base code works when there are no NaN
s:
l = [[1], [2, 3, 4], [5]] [item for sublist in l for item in sublist] > [1, 2, 3, 4, 5]
Now if I have the following case it will break:
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
which makes sense, but I need this border case to be handled and I’m not sure how to add the special case condition in the list comprehension above. How can I modify such comprehension to cover such case?
The expected output is:
[1, 2, 3, 4, np.nan, 5]
Advertisement
Answer
I would check to see if the sublist is iterable. If not then wrap it in a tuple to then be flattened:
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]
Using chain
from itertools
would make it a bit cleaner:
from itertools import chain list(chain(*(sublist if hasattr(sublist, "__iter__") else (sublist, ) for sublist in l))) >>> [1, 2, 3, 4, nan, 5]