Is there a way to write successive for loops in one line in Python?

For example, I would like to create a set with n 0s and m 1s, [0, 0, ...0, 1, 1, ..., 1]. Is there a way to do something like [0 for i in range(n), 1 for j in range(m)]?

Solution 1:

You can create the two lists separately, then add them like this:

[0 for i in range(n)] + [1 for j in range(m)]

Or simply use list multiplication:

[0]*n + [1]*m

Solution 2:

In the event you wanted to do this without having to add two lists together (e.g. if n and m are very large and you want to be able to generate this sequence without putting it all into a list), you could iterate over the combined range with a conditional:

>>> n, m = 2, 3
>>> [0 if i < n else 1 for i in range(n + m)]
[0, 0, 1, 1, 1]

You can also use itertools.chain to "add" two generators together:

>>> from itertools import chain
>>> list(chain((0 for _ in range(n)), (1 for _ in range(m))))
[0, 0, 1, 1, 1]

Solution 3:

Stop thinking in terms of doing things in a single line. Think in terms of using useful, modular, extensible abstractions. For these sorts of things, itertools provides a lot of useful abstractions.

>>> from itertools import repeat, chain
>>> list(chain(repeat(0, 4), repeat(1, 3)))
[0, 0, 0, 0, 1, 1, 1]

Start building your own useful abstractions by combining smaller useful abstractions:

>>> from itertools import repeat, chain, starmap
>>> def builder(*pairs, consume=list):
...     return consume(chain.from_iterable(starmap(repeat, pairs)))
...
>>> builder((0,3), (1, 4), (2,1), (3, 6))
[0, 0, 0, 1, 1, 1, 1, 2, 3, 3, 3, 3, 3, 3]

Although, sometimes, simpler is better, and learn to take advantage of the various syntactic sugars provided by the language:

>>> [*repeat(0, 4), *repeat(1, 3), *repeat(2,2)]
[0, 0, 0, 0, 1, 1, 1, 2, 2]