How to split a list into sublists based on a separator, similar to str.split()?

Question

Given a list like:

[a, SEP, b, c, SEP, SEP, d]

how do I split it into a list of sublists:

[[a], [b, c], [], [d]]

Effectively I need an equivalent of str.split() for lists. I can hack together something, but I can't seem to be able to come up with anything neat and/or pythonic.

I get the input from an iterator, so a generator working on that is acceptable as well.

More examples:

[a, SEP, SEP, SEP] -> [[a], [], [], []]

[a, b, c] -> [[a, b, c]]

[SEP] -> [[], []]

Answer 1

A simple generator will work for all of the cases in your question:

def split(sequence, sep):
    chunk = []
    for val in sequence:
        if val == sep:
            yield chunk
            chunk = []
        else:
            chunk.append(val)
    yield chunk

Answer 2

My first ever Python program :)

from pprint import pprint
my_array = ["a", "SEP", "SEP", "SEP"]
my_temp = []
my_final = []
for item in my_array:
  if item != "SEP":
    my_temp.append(item)
  else:
    my_final.append(my_temp);
    my_temp = []
pprint(my_final);

Answer 3

I am not sure if there's an easy itertools.groupby solution, but here is an iterative approach that should work:

def mySplit(iterable, sep):
    output = []
    sepcount = 0
    current_output = []
    for i, elem in enumerate(iterable):
        if elem != sep:
            sepcount = 0
            current_output.append(elem)
            if (i==(len(iterable)-1)):
                output.append(current_output)
        else:
            if current_output: 
                output.append(current_output)
                current_output = []

            sepcount+=1

            if (i==0) or (sepcount > 1):
                output.append([])
            if (i==(len(iterable)-1)):
                output.append([])

    return output

Testing on your examples:

testLists = [
    ['a', 'SEP', 'b', 'c', 'SEP', 'SEP', 'd'],
    ["a", "SEP", "SEP", "SEP"],
    ["SEP"],
    ["a", "b", "c"]
]

for tl in testLists:
    print(mySplit(tl, sep="SEP"))
#[['a'], ['b', 'c'], [], ['d']]
#[['a'], [], [], []]
#[[], []]
#[['a', 'b', 'c']]

This is analogous to the result you would get if examples were actually strings and you used str.split(sep):

for tl in testLists:
    print("".join(tl).split("SEP"))
#['a', 'bc', '', 'd']
#['a', '', '', '']
#['', '']
#['abc']

By the way, if the elements in your lists were always guaranteed to be strings, you could simply do:

for tl in testLists:
    print([list(x) for x in "".join(tl).split("SEP")])
#[['a'], ['b', 'c'], [], ['d']]
#[['a'], [], [], []]
#[[], []]
#[['a', 'b', 'c']]

But the mySplit() function is more general.

Answer 4

For list or tuple objects you can use the following:

def split(seq, sep):
    start, stop = 0, -1
    while start < len(seq):
        try:
            stop = seq.index(sep, start)
        except ValueError:
            yield seq[start:]
            break
        yield seq[start:stop]
        start = stop + 1
    else:
        if stop == len(seq) - 1:
            yield []

I won't work with a generator but it's fast.

Answer 5

You can use itertools.takewhile:

def split(seq, sep):
    seq, peek = iter(seq), sep
    while True:
        try:
            peek = next(seq)
        except StopIteration:
            break
        yield list(it.takewhile(sep.__ne__, it.chain((peek,), seq)))
    if peek == sep:
        yield []

The it.chain part is to find out when the seq is exhausted. Note that with this approach it's easy to yield generators instead of lists if desired.

Answer 6

If you prefer a list comprehension, then you can resort to filtering indices and slicing, using itertools.pairwise:

seq = ['a', 'SEP', 'b', 'c', 'SEP', 'SEP', 'd']
[seq[a + 1 : b]
 for (a, b) in itertools.pairwise(
     [-1] + [i for i in range(len(seq)) if seq[i] == 'SEP'] + [len(seq)])]

→

[['a'], ['b', 'c'], [], ['d']]

Answer 7

I would define the following function to solve that problem.

l = ['a', 'SEP', 'b', 'c', 'SEP', 'SEP', 'd']

def sublist_with_words(word, search_list):
    res = []
    for i in range(search_list.count(word)):
        index = search_list.index(word)
        res.append(search_list[:index])
        search_list = search_list[index+1:]
    res.append(search_list)
    return res

When I try the cases you gave:

print(sublist_with_words(word = 'SEP', search_list=l))
print(sublist_with_words(word = 'SEP', search_list=['a', 'b', 'c']))
print(sublist_with_words(word = 'SEP', search_list=['SEP']))

The output is:

[['a'], ['b', 'c'], [], ['d']]
[['a', 'b', 'c']]
[[], []]

Answer 8

itertools.takewhile @a_guest's approach simplified:

def split(seq, sep):
    from itertools import takewhile
    iterator = iter(seq)
    while subseq := list(takewhile(lambda x: x != sep, iterator)):
        yield subseq

Please note it returns on first empty subsequence.

Answer 9

The following is a non-generic solution that (most probably) only works on list of ints:

import re

def split_list(nums, n):
    nums_str = str(nums)
    splits = nums_str.split(f"{n},")

    patc = re.compile(r"\d+")
    group = []
    for part in splits:
        group.append([int(v) for v in patc.findall(part)])

    return group

if __name__ == "__main__":
    l = [1, 2, 3, 4, 3, 6, 7, 3, 8, 9, 10]
    n = 3
    split_l = split_list(l, n)
    assert split_l == [[1, 2], [4], [6, 7], [8, 9, 10]]