Using itertools groupby to create a list of lists

Question

I'm getting a list of items in the format parent.id_child.id, like 1_2. I tried to group the child's id by parent id, for example: from the input ['1_2', '2_2', '1_1', '2_1', '1_3'] I need the output [['1','2','3'], ['1','2']]. I have tried this:

inputlist = ['1_2', '1_1', '2_1', '1_3', '2_2' ]
outputlist= [item.split('_') for item in inputlist]
outputlist.sort()
final = [list(group) for key, group in itertools.groupby(outputlist, lambda x: x[0])]

This groups the elements correctly, by I need to obtain only the second element of each item. How can I achieve this? Also, can I do the whole thing in a single list comprehension sentence?

Answer 1

Use a list comprehension, yes; the values themselves are passed into each group iterator unchanged so you need to select again:

final = [[g[1] for g in group] for key, group in itertools.groupby(outputlist, lambda x: x[0])]

You can do the whole thing in a single expression by nesting the splitting into the groupby call, but this becomes rather ugly fast, even when split across multiple lines:

final = [
    [g[1] for g in group]
    for key, group in itertools.groupby(
        sorted(item.split('_') for item in inputlist),
        lambda x: x[0])]

You could avoid sorting the whole input list and only sort the smaller groups by using a dictionary to do the grouping. Dependending on the size of your ids, you may want to sort your ids numerically as well (since text sorting is done lexicographically):

per_parent = {}
for item in inputlist:
    parent, child = item.split('_', 1)
    per_parent.setdefault(parent, []).append(child)
final = [children for parent, children in sorted(
    per_parent.items(), key=lambda pc: int(pc[0]))]

In Python 2, use iteritems() rather than items() to avoid building an intermediate list.