Itertools Recipe: Round Robin

The itertools docs has a ton of slick recipes for using the library to good effect. Some of the code is more useful than illustrative, so I wanted to use these notebooks to break down a few of the functions.

This is

# poor import style, but I want to copy-paste the code
# as-is from the docs

from itertools import *
import itertools

roundrobin()

def roundrobin(*iterables):
    "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
    # Recipe credited to George Sakkis
    num_active = len(iterables)
    nexts = cycle(iter(it).__next__ for it in iterables)
    while num_active:
        try:
            for next in nexts:
                yield next()
        except StopIteration:
            # Remove the iterator we just exhausted from the cycle.
            num_active -= 1
            nexts = cycle(islice(nexts, num_active))

Demo

The second value should be sequential, b dropping out, and then a

player_a = ['a1', 'a2', 'a3']
player_b = ['b1', 'b2']
player_c = ['c1', 'c2', 'c3', 'c4']

for item in roundrobin(player_a, player_b, player_c):
    print(item, end=' ')

a1 b1 c1 a2 b2 c2 a3 c3 c4 

Mississippi… sort of

m = 'm'
i = 'iii'
s = 'ssss'
p = 'pp'

for item in roundrobin(m, i, s, p):
    print(item, end='')

mispispiss

lol

Why this works

This one relies heavily on some Functional Programming Magic™, but beyond the main trick, the function is easy enough to understand.

The nexts iterable

The function starts off by creating a straight-forward num_active variable that tracks how many iterables we’re juggling.

The nexts line merits some unpacking:

• At the outer level, it invokes the itertools.cycle() method to repeatedly yield the next result of each iterable
• We wrap each value it in the iter(), on the off chance one of *iterables was a standalone string or something
• Finally magic part, outlined above is in the fact that we’re cycling through __next__() functions, not the values that they yield

From here, we execute the loop while there are still values to yield in any of the iterables

Using and redefining nexts

Because it’s built using itertools.cycle(), the for next in nexts: portion will loop indefinitely, yielding the next value of each iterable, until the first time that it exhausts one of them, which is why we’ve got it stuffed inside the try block

At that point, we catch the StopIteration exception, decrement our number of active iterables by one, then do some more clever Functional Programming Magic™.

Revisiting the player_a, player_b, player_c example from above, after looping a couple times, we’ve yielded a1 b1 c1 a2 b2 c2 a3, pushed player_a to the back of the cycle, and are getting ready to serve the next value of player_b

B: 
C: c3, c4
A: 

However, player_b’s out of values and is about to kick us the StopIteration exception

C: c3, c4
A: 
B: StopIteration

But when it does that, it, critically, gets moved to the back of the cycle

At that point, we build a new cycle using itertools.islice(). We still use the nexts iterable, but because we decremented num_activive in the previous line, this means that we’re cycling through only the first n-1 iterables in nexts.

Which, in this case, is player_a and player_c– we’ve discarded the empty player_b from the rotation