zhaodaolimeng
zhaodaolimeng

Reputation: 626

python repeat list elements in an iterator

Is there any way to create an iterator to repeat elements in a list certain times? For example, a list is given:

color = ['r', 'g', 'b']

Is there a way to create a iterator in form of itertools.repeatlist(color, 7) that can produce the following list?

color_list = ['r', 'g', 'b', 'r', 'g', 'b', 'r']

Upvotes: 7

Views: 5170

Answers (2)

Martijn Pieters
Martijn Pieters

Reputation: 1122222

You can use itertools.cycle() together with itertools.islice() to build your repeatlist() function:

from itertools import cycle, islice

def repeatlist(it, count):
    return islice(cycle(it), count)

This returns a new iterator; call list() on it if you must have a list object.

Demo:

>>> from itertools import cycle, islice
>>> def repeatlist(it, count):
...     return islice(cycle(it), count)
...
>>> color = ['r', 'g', 'b']
>>> list(repeatlist(color, 7))
['r', 'g', 'b', 'r', 'g', 'b', 'r']

Upvotes: 18

youkaichao
youkaichao

Reputation: 2364

The documentation of cycle says:

Note, this member of the toolkit may require significant auxiliary storage (depending on the length of the iterable).

I'm curious why python doesn't provide a more efficient implementation:

def cycle(it):
    while True:
        for x in it:
            yield x

def repeatlist(it, count):
    return [x for (i, x) in zip(range(count), cycle(it))]

This way, you don't need to save a whole copy of the list. And it works if list is infinitely long.

Upvotes: 1

Related Questions