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 in [b] to produce this:

[0, 5, 5, 1]

ie 0 occurs 1 time, 5 occurs 2 times, and 1 occurs 1 time.

Answer 1

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 2

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

Answer 3

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 4

Usinglist comprehension , what you want to obtain is:

[a[0]]*b[0] + [a[1]]*b[1] + [a[2]]*b[2]

which, for short lists, is convenient to achieve using sum(list, [])

What does the built-in function sum do with sum(list, [])?

as it might be even faster (and more straightforward) than itertools.chain + itertools.repeat for short lists.

sum([ [a_i]*b_i for a_i, b_i in zip(a, b) ], [])

For long lists, stick with itertools.chain and itertools.repeat .