Skip to content
Advertisement

Split list at a specific value

I am trying to write a code that splits lists in a class of lists in two when a certain value is a middle element of the list and then produce two lists where the middle element becomes the end element in the first list and the first element in the second one.

There can be more than n middle elements in the list so the result must be n+1 lists.

Example:

A = [[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15],[16,17,18,19,20,21,22,23,24,25],[26,27,28,29]]

P = [4,7,13,20]
n = len(Points) # in this case n = 4

I am looking for a result that looks like this:

A = [[0,1,2,3,4],[4,5,6,7],[7,8,9,10,11,12,13],[13,14,15],[16,17,18,19,20],[20,21,22,23,24,25],[26,27,28,29]]

Since n = 4 and it will produce 5 lists, note that the answer has 6 lists because the last list doesn’t have any value of P in and therefore stays intact.

I haven’t been able to produce anything as I am new to python and it is hard to formulate this problem.

Any help is appreciated!

Advertisement

Answer

You can first recover all indices of the provided values and then slice accordingly.

Code

def split_at_values(lst, values):
    indices = [i for i, x in enumerate(lst) if x in values]
    for start, end in zip([0, *indices], [*indices, len(lst)]):
        yield lst[start:end+1]
        # Note: remove +1 for separator to only appear in right side slice

Example

values =  {4, 7, 13, 20}
lst = [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]

print(*split_at_values(lst, values))

Output

[0, 1, 2, 3, 4] [4, 5, 6, 7] [7, 8, 9, 10, 11, 12, 13] [13, 14, 15]

You can then apply this iteratively to you input list A to get the desired result. Alternatively you can use itertools.chain.from_iterable.

from itertools import chain

values = {4, 7, 13, 20}
lst_A = [[0, 1, 2, 3, 4, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
         [16, 17, 18, 19, 20, 21, 22, 23, 24, 25],
         [26, 27, 28, 29]]

output = list(chain.from_iterable(split_at_values(sublst, values) for sublst in lst_A))

print(output)

Output

[[0, 1, 2, 3, 4],
 [4, 5, 6, 7],
 [7, 8, 9, 10, 11, 12, 13],
 [13, 14, 15],
 [16, 17, 18, 19, 20],
 [20, 21, 22, 23, 24, 25],
 [26, 27, 28, 29]]
Advertisement