3

I would like to flatten a list but keep NaNs. The following base code works when there are no NaNs:

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]
SkyWalker
  • 13,729
  • 18
  • 91
  • 187

2 Answers2

3

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]

Source for checking if iterable (StackOverflow)

Mandera
  • 2,647
  • 3
  • 21
  • 26
1

Trying different things I found a possible solution:

[item for sublist in l for item in (sublist if type(sublist)==list else [sublist])]
> [1, 2, 3, 4, np.nan, 5]
SkyWalker
  • 13,729
  • 18
  • 91
  • 187
  • Typechecking is never a good approach, you can instead use `isinstance` builtin, more details at: https://stackoverflow.com/questions/17493159/should-i-force-python-type-checking – ThePyGuy Aug 17 '22 at 10:38