按特定值将列表拆分为列表

Question

I have a list:

['S1', 'S2', 'S6', 'S1', 'S2', 'S3', 'S4', 'S5', 'S1', 'S2', 'S5', 'S1',
 'S2', 'S4', 'S5', 'S1', 'S2', 'S4', 'S5', 'S1', 'S2', 'S3', 'S6']

and I want to split by next S1:

[['S1', 'S2', 'S6']['S1', 'S2', 'S3', 'S4', 'S5'],['S1', 'S2', 'S4', 'S5]...]

My code is:

size = len(steps)
idx_list = [idx + 1 for idx, val in
            enumerate(steps) if val == 'S1'] 


res = [steps[i: j] for i, j in
        zip([0] + idx_list, idx_list + 
        ([size] if idx_list[-1] != size else []))] 

print("The list after splitting by a value : " + str(res))

It splits the list as:

[['S1'], ['S2', 'S6', 'S1'], ['S2', 'S3', 'S4', 'S5', 'S1'], 
 ['S2', 'S5', 'S1'], ['S2', 'S4', 'S5', 'S1'], ['S2', 'S4', 'S5', 'S1']..

Can you please help to rectify it!

Answer 1

You can use itertools.groupby :

from itertools import groupby

lst = ['S1', 'S2', 'S6', 'S1', 'S2', 'S3', 'S4', 'S5', 'S1', 'S2', 'S5', 'S1', 'S2', 'S4', 'S5', 'S1', 'S2', 'S4', 'S5', 'S1', 'S2', 'S3', 'S6']

splitby = 'S1'
res = [[splitby] + list(g) for k, g in groupby(lst, key=lambda x: x != splitby) if k]

# [['S1', 'S2', 'S6'], ['S1', 'S2', 'S3', 'S4', 'S5'], ['S1', 'S2', 'S5'], ['S1', 'S2', 'S4', 'S5'], ['S1', 'S2', 'S4', 'S5'], ['S1', 'S2', 'S3', 'S6']]

Answer 2

You have an off-by-one error . Change the following line:

idx_list = [idx + 1 for idx, val in
            enumerate(steps) if val == 'S1'] 

to

idx_list = [idx for idx, val in
            enumerate(steps) if val == 'S1' and idx > 0] 

The result should then be:

[['S1', 'S2', 'S6'], ['S1', 'S2', 'S3', 'S4', 'S5'], 
 ['S1', 'S2', 'S5'], ['S1', 'S2', 'S4', 'S5'], 
 ['S1', 'S2', 'S4', 'S5'], ['S1', 'S2', 'S3', 'S6']]