Repeat elements in one list based on elements from another

Question

Given the following lists:

a = [0, 5, 1]
b = [1, 2, 1]

I\'d like to repeat each element of [a] by the number of its corresponding position

Answer 1

This can be done straightforwardly using enumerate():

a = [0, 5, 1]
b = [1, 2, 1]
[ele for i, ele in enumerate(a) for j in range(b[i])]

Answer 2

In [7]: a = [0, 5, 1]

In [8]: b = [1, 2, 1]

In [9]: list(itertools.chain(*(itertools.repeat(elem, n) for elem, n in zip(a, b))))
Out[9]: [0, 5, 5, 1]

In [10]: b = [2, 3, 4]

In [11]: list(itertools.chain(*(itertools.repeat(elem, n) for elem, n in zip(a, b))))
Out[11]: [0, 0, 5, 5, 5, 1, 1, 1, 1]

The pieces here are as follows:

• itertools.repeat(elem, n) - repeat elem n times
• zip(a, b) Make a list of 2-tuples out of the two lists, pairing each element with the corresponding element in the other list. This gives you exactly what you need to pass to itertools.repeat in your use case.
• itertools.chain - flattens the resulting list of iterators into a single list of values. You can either chain(*iterable) as I have done or chain.from_iterable(iterable) as Martijn Peters does.

Answer 3

Use the zip() function with itertools.repeat() and itertools.chain.from_iterable():

try:
    # use iterator zip on Python 2 too
    from future_builtins import zip
except ImportError:
    pass
from itertools import repeat, chain

list(chain.from_iterable(repeat(value, count) for value, count in zip(a, b)))

(I added a Python 2 forward-compatible import for those that can't switch to Python 3 yet).

Demo:

>>> from itertools import repeat, chain
>>> a = [0, 5, 1]
>>> b = [1, 2, 1]
>>> list(chain.from_iterable(repeat(value, count) for value, count in zip(a, b)))
[0, 5, 5, 1]

An alternative approach would be to use a list comprehension; this is slower as repeating elements is done in bytecode instead of C:

[value for value, count in zip(a, b) for _ in range(count)]