How To Split A List Into Subsets Based On A Pattern?

I'm doing this but it feels this can be achieved with much less code. It is Python after all. Starting with a list, I split that list into subsets based on a string prefix. # Split

Solution 1:

You could use itertools.groupby:

>>>import itertools>>>mylist = ['sub_0_a', 'sub_0_b', 'sub_1_a', 'sub_1_b']>>>for k,v in itertools.groupby(mylist,key=lambda x:x[:5]):...print k, list(v)... 
sub_0 ['sub_0_a', 'sub_0_b']
sub_1 ['sub_1_a', 'sub_1_b']

or exactly as you specified it:

>>> [list(v) for k,v in itertools.groupby(mylist,key=lambda x:x[:5])]
[['sub_0_a', 'sub_0_b'], ['sub_1_a', 'sub_1_b']]

Of course, the common caveats apply (Make sure your list is sorted with the same key you're using to group), and you might need a slightly more complicated key function for real world data...

Solution 2:

In [28]: mylist = ['sub_0_a', 'sub_0_b', 'sub_1_a', 'sub_1_b']

In [29]: lis=[]

In [30]: for x in mylist:
    i=x.split("_")[1]
    try:
        lis[int(i)].append(x)
    except:    
        lis.append([])
        lis[-1].append(x)
   ....:         

In [31]: lis
Out[31]: [['sub_0_a', 'sub_0_b'], ['sub_1_a', 'sub_1_b']]

Solution 3:

Use itertools' groupby:

def get_field_sub(x): return x.split('_')[1]

mylist = sorted(mylist, key=get_field_sub)
[ (x, list(y)) for x, y in groupby(mylist, get_field_sub)]