Flattening lists - sometimes you need better data

When you have a problem that’s hard to solve, you can make it easier by either

• changing the code, or
• changing the data

How to flatten a list

The question of flattening lists seems to come up a lot in Python discussion forums.
You have a list of lists like this

[ [1, 2], [3, 4], [5, 6] ]

and you want to turn it into

[1, 2, 3, 4, 5, 6]

itertools to the rescue

You probably know about the itertools.chain.from_iterable function which is designed to solve this exact problem.

a = [ [1, 2], [3, 4], [5, 6] ]

import itertools
flattened_a = list(itertools.chain.from_iterable(a))

But there are times when this will fail. For example

a = [ [1, 2], 3, 4, 5, 6 ]

As the name says, from_iterable expects all the list members to be iterable objects. This example fails because integers are not iterable. What should you do?

• change your algorithm?
• change your data?

Change your data

The from_iterable function would work fine if our data was [ [1, 2], [3], [4], [5], [6] ]. So instead of changing the code, why not change the data?

a = [ [1, 2], 3, 4, 5, 6 ]
list_of_lists = [i if isinstance(i, list) else [i] for i in a]
itertools.chain.from_iterable(list_of_lists)

Anything that’s not a list gets wrapped in square brackets. Now you have a list of lists and from_iterables works as expected.

Don’t get trapped into a single point of view

Don’t forget that as a programmer you have control over both the code and the data.

Did you enjoy this post? I have a book.

​