10

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.

Georgy
  • 12,464
  • 7
  • 65
  • 73
Dance Party2
  • 7,214
  • 17
  • 59
  • 106
  • 2
    It looks like you want us to write some code for you. While many users are willing to produce code for a coder in distress, they usually only help when the poster has already tried to solve the problem on their own. A good way to demonstrate this effort is to include the code you've written so far, example input (if there is any), the expected output, and the output you actually get (console output, stack traces, compiler errors - whatever is applicable). The more detail you provide, the more answers you are likely to receive. – Martijn Pieters Jun 02 '16 at 15:02
  • Does this answer your question? [Repeat each item in a list a number of times specified in another list](https://stackoverflow.com/questions/33382474/repeat-each-item-in-a-list-a-number-of-times-specified-in-another-list) – Georgy Aug 26 '20 at 16:28

4 Answers4

9
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.
Two-Bit Alchemist
  • 17,966
  • 6
  • 47
  • 82
7

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)]
Martijn Pieters
  • 1,048,767
  • 296
  • 4,058
  • 3,343
4

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])]
mattinthehat
  • 51
  • 2
0

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.

Pietro D'Antuono
  • 352
  • 1
  • 11