fishman11235
fishman11235

Reputation: 11

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)]?

Upvotes: 1

Views: 80

Answers (5)

juanpa.arrivillaga
juanpa.arrivillaga

Reputation: 96172

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]

Upvotes: 2

Simon Hawe
Simon Hawe

Reputation: 4539

Yet another potential solution for your problem using list comprehension and 2 for loops

n,m = 3,4
[i for i,n_numbers in enumerate((n,m)) for _ in range(n_numbers) ]

Upvotes: 0

Samwise
Samwise

Reputation: 71517

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]

Upvotes: 2

XxJames07-
XxJames07-

Reputation: 1826

you could do:

lst = [0]*n+[1]*m

Upvotes: 1

Djaouad
Djaouad

Reputation: 22776

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

Upvotes: 4

Related Questions