How to "round-robin" a single iterable based on some criteria?

Question

I have a list of dictionaries, which I want to make round robin sorted.

sample = [
    {'source': 'G', '"serial"': '0'},
    {'source': 'G', '"serial"': '1'},
    {'source': 'G', '"serial"': '2'},
    {'source': 'P', '"serial"': '30'},
    {'source': 'P', '"serial"': '0'},
    {'source': 'P', '"serial"': '1'},
    {'source': 'P', '"serial"': '2'},
    {'source': 'P', '"serial"': '3'},
    {'source': 'T', '"serial"': '2'},
    {'source': 'T', '"serial"': '3'}
]

I want this result:

sample_solved = [
    {'source': 'G', '"serial"': '0'},
    {'source': 'P', '"serial"': '30'},
    {'source': 'T', '"serial"': '2'},
    {'source': 'G', '"serial"': '1'},
    {'source': 'P', '"serial"': '1'},
    {'source': 'T', '"serial"': '3'},
    {'source': 'G', '"serial"': '2'},
    {'source': 'P', '"serial"': '0'},
    {'source': 'P', '"serial"': '2'},
    {'source': 'P', '"serial"': '3'}
]

The way I solved it is as follows:

def roundrobin(*iterables):
    # took from here https://docs.python.org/3/library/itertools.html#itertools-recipes

    "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
    # Recipe credited to George Sakkis

    pending = len(iterables)
    nexts = cycle(iter(it).__next__ for it in iterables)
    while pending:
        try:
            for next in nexts:
                yield next()
        except StopIteration:
            pending -= 1
            nexts = cycle(islice(nexts, pending))

def solve():
    items_by_sources = collections.defaultdict(list)

    for item in sample2:
         items_by_sources[item["source"]].append(item)

    t, p, g = items_by_sources.values()

    print(list(roundrobin(t, p, g)))

Using Python's defaultdict to separate the items by source and then using the roundrobin solution which I got from Python's docs.

But, the solution, does not cover all the cases, for example t, p, g = items_by_sources.values() will break when one source is missing or a new source has been added.

How can I make a solution to cover more edge cases and make the solution pythonic?

Answer 1

Here's a solution using itertools.groupby() to split your input into the appropriate groups:

from itertools import groupby

def grouprobin(iterable, key):
    groups = [list(g) for k, g in groupby(iterable, key)]
    while groups:
        group = groups.pop(0)
        yield group.pop(0)
        if group:
            groups.append(group)

Because of the way groupby() works, the clever use of iterators in the version of roundrobin() you took from the docs isn't very helpful, so I've rewritten it in a way that's hopefully easier to follow:

1. Group the iterable by key

2. While you still have any groups left:

   1. Pop the first group from the front of the list of groups

   2. Pop the first item from that group, and yield it.

   3. If there are still items in the group, append it back to the end of the list.

Here it is in action:

>>> sample_solved = list(grouprobin(sample, key=lambda d: d['source']))
>>> from pprint import pprint
>>> pprint(sample_solved)
[{'"serial"': '0', 'source': 'G'},
 {'"serial"': '30', 'source': 'P'},
 {'"serial"': '2', 'source': 'T'},
 {'"serial"': '1', 'source': 'G'},
 {'"serial"': '0', 'source': 'P'},
 {'"serial"': '3', 'source': 'T'},
 {'"serial"': '2', 'source': 'G'},
 {'"serial"': '1', 'source': 'P'},
 {'"serial"': '2', 'source': 'P'},
 {'"serial"': '3', 'source': 'P'}]

The version of grouprobin() above assumes that your list is already sorted. If not, it'll need to be sorted before it's grouped:

def grouprobin(iterable, key):
    groups = [list(g) for k, g in groupby(sorted(iterable, key=key), key)]
    while groups:
        group = groups.pop(0)
        yield group.pop(0)
        if group:
            groups.append(group)