Pythonic way to split a list after elements for which a given predicate is true

Question

Assume you have a list of arbitrary elements like

['monkey', 'deer', 'lion', 'giraffe', 'lion', 'eagle', 'lion', 'fish']

which should be split into sublists after each element for which a given predicate, e.g.

is_lion(element)

returns True. The above example should become

[['monkey', 'deer', 'lion'], ['giraffe', 'lion'], ['eagle', 'lion'], ['fish']]

Is there a pythonic way of doing it?

Answer 1

The easiest way is probably:

out = [[]]
for element in lst:
    out[-1].append(element)
    if predicate(element):
        out.append([])

Note that this would leave an empty list at the end of out, if predicate(element): for the last element. You can remove this by adding:

out = [l for l in out if l]

Answer 2

Just because we can, a functional one-liner:

from functools import reduce

reduce(lambda out, x: out[:-1] + [out[-1] + [x]] if not predicate(x) else out + [[x]], x, [[]])

Answer 3

I rather like this solution:

def f(outs, x):
    if outs[-1][-1:] == ["lion"]:
        outs.append([])
    outs[-1].append(x)
    return outs

def splitAfterLion(xs):
    return reduce(f,xs,[[]])

It might not be very pythonic, more functional. But it's short and does not suffer from trailing empty lists in the result.

Answer 4

>>> import itertools
>>> l = ['monkey', 'deer', 'lion', 'giraffe', 'lion', 'eagle', 'lion', 'fish']
>>> f = lambda i: i == "lion"
>>> a = [list(j) for i, j in itertools.groupby(l, f)]
>>> a
[['monkey', 'deer'], ['lion'], ['giraffe'], ['lion'], ['eagle'], ['lion'], ['fish']]
>>> [i+j for i, j in zip(a[::2], a[1::2])]
[['monkey', 'deer', 'lion'], ['giraffe', 'lion'], ['eagle', 'lion']]

Edit:

>>> [i+j for i, j in itertools.zip_longest(a[::2], a[1::2], fillvalue=[])]
[['monkey', 'deer', 'lion'], ['giraffe', 'lion'], ['eagle', 'lion'], ['fish']]

Answer 5

Just another way of doing it by getting the index without using itertool, please let me know if that works for you:

#!/usr/bin/python

ls = ['monkey', 'deer', 'lion', 'giraffe', 'lion', 'eagle', 'lion', 'fish', 'fish']

def is_lion(elm):
    return elm in ls

def mark_it(nm):
    ind = [ x+1 for x,y in enumerate(ls) if y == nm ]
    if ind[-1] < len(ls):
        ind.append(len(ls))
    return ind

def merge_it(ind):
    return [list(ls[x[0]:x[1]]) for x in zip(ind[::], ind[1::])]

name = 'lion'
if is_lion(name):
    index = [0]
    index.extend(mark_it(name))
    print merge_it(index)
else:
    print 'not found'

Output:

[['monkey', 'deer', 'lion'], ['giraffe', 'lion'], ['eagle', 'lion'], ['fish', 'fish']]

Answer 6

Here is a solution:

def is_lion(a, element):
    start = 0
    for key,value in enumerate(a):
        if value == element:
            yield a[start:key+1]
            start = key+1

    # print out the last sub-list
    if value != 'lion':
        yield a[start:key+1]


a = ['monkey', 'deer', 'lion', 'giraffe', 'lion', 'eagle', 'lion', 'fish']

print [x for x in is_lion(a, 'lion')]