CODEWITHSUNDEEP
pythonrepeat
Dance Party2
Dance Party2

Reputation: 7536

Repeat elements in one list based on elements from another

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]

i.e. 0 occurs 1 time, 5 occurs 2 times, and 1 occurs 1 time.

Upvotes: 10

Views: 2072

Answers (4)

Pietro D'Antuono
Pietro D'Antuono

Reputation: 412

Using list 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.

Upvotes: 0

mattinthehat
mattinthehat

Reputation: 61

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

Upvotes: 4

Martijn Pieters
Martijn Pieters

Reputation: 1121724

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

Upvotes: 7

Two-Bit Alchemist
Two-Bit Alchemist

Reputation: 18457

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.

Upvotes: 9

Related Questions