Divide a list in python to sub lists by condition-CodePudding

So I have a list of variables name, for example like this:

['aS3',
 'a1',
 'a3',
 'a3_cod_1',
 'a3_cod_2',
 'a3_cod_3',
 'a3_cod_4',
 'a3_cod_5',
 'a3_cod_6',
 'a12A_1',
 'a12A_2',
 'a12A_3',
 'a12A_4',
 'a12A_5',
 'a12A_6',
 'a12A_7',
 'a12A_8',
 'a12A_9',
 'a12A_10',
 'aD7A',
 'aD7B',]

So I need to get from it lists of variables that starts with the same symbols before first underscore. For this example we should have something like this:

['a3', 'a3_cod_1', 'a3_cod_2', 'a3_cod_3', 'a3_cod_4', 'a3_cod_5', 'a3_cod_6']

and

['a12A_1', 'a12A_2', 'a12A_3', 'a12A_4', 'a12A_5', 'a12A_6', 'a12A_7', 'a12A_8', 'a12A_9', 'a12A_10']

I tried to do it with some for and if's statements but it looks real heavy. So I want to find more elegant solution

CodePudding user response：

Use a dictionary to categorize the elements by their prefix.

collections = {}
for element in lst:                       # lst is the full list
    prefix = element.split("_")[0]        # get the prefix
    collections.setdefault(prefix,[])     # set the key, value
    collections[prefix].append(element)   # append the element to the value
print(list(collections.values()))         # output the values as a list

Output:

[
['aS3'], 
['a1'], 
['a3', 'a3_cod_1', 'a3_cod_2', 'a3_cod_3', 'a3_cod_4', 'a3_cod_5', 'a3_cod_6'], 
['a12A_1', 'a12A_2', 'a12A_3', 'a12A_4', 'a12A_5', 'a12A_6', 'a12A_7', 'a12A_8', 'a12A_9', 'a12A_10'], 
['aD7A'], ['aD7B']
]

CodePudding user response：

prefixes = set([f"{x.split('_')[0]}_" for x in a if "_" in x]) 

d = {}
for p in prefixes:
    d[p] = [x for x in a if x.startswith(p)]


{'a3_': ['a3_cod_1', 'a3_cod_2', 'a3_cod_3', 'a3_cod_4', 'a3_cod_5', 'a3_cod_6'],
 'a12A_': ['a12A_1', 'a12A_2', 'a12A_3', 'a12A_4', 'a12A_5', 'a12A_6', 'a12A_7', 'a12A_8', 'a12A_9', 'a12A_10']}