How to get sublist from list with keyword for start and end multiple times

Viewed 221

I have a list for example ['tree', 'water', 'dog', 'soap', 'cat', 'bird', 'tree', 'dog', 'water'] I want a function that takes in any list with a start word and ends word and gives me back sub-lists of them. I want it to search the entire list for the start and end word. So if my start word is tree and end word is water it would give me back two lists consisting of ['tree', water'] and ['tree', 'dog', 'water']. And if there is no word tree or water then the function should just skip it.

I've tried doing

def sublist(words, start, end):
    start_index = words.index(start)
    end_index = words.index(end)
    sublist = words[start_index:end_index+1]

words = ['tree', 'water', 'dog', 'soap', 'cat', 'bird', 'tree', 'dog', 'water']
start = 'tree'
end = 'water'
sublist(words, start, end)

but this only gives back one list of ['tree', 'water']. I want ['tree,' water],['tree', 'dog', 'water'], I don't know how to keep continuing after I find the first segment. It would also error if the start or end word is not on the list.

2 Answers

You could keep indexing until failure:

def sublists(words, start, end):
    end_index = -1
    try:
        while True:
            start_index = words.index(start, end_index + 1)
            end_index = words.index(end, start_index + 1)
            yield words[start_index : end_index+1]
    except ValueError:
        pass

Or a regex way, first turning the list into a string like 'se----s-e':

import re

def sublists(words, start, end):
    s = ''.join('s' if w == start else 'e' if w == end else '-' for w in words)
    for match in re.finditer('s-*e', s):
        yield words[slice(*match.span())]

Or a somewhat abusive way:

def sublists(words, start, end):
    it = iter(words)
    while start in it:
        sub = [start]
        if end in (sub.append(word) or word for word in it):
            yield sub

Not sure what to think of this version of it:

from itertools import filterfalse

def sublists(words, start, end):
    it = iter(words)
    while start in it:
        sub = [start]
        if end in filterfalse(sub.append, it):
            yield sub

Non-abusive version of it:

def sublists(words, start, end):
    it = iter(words)
    while start in it:
        sub = [start]
        for word in it:
            sub.append(word)
            if word == end:
                yield sub
                break

maybe this can deal with it:

def sublist(words, start, end):
    sublists = []                           # prepare sublists for return
    sublist = []                            # initialze sublist
    for word in words:                      # for each word
        if len(sublist):                    # sublist already has word(s)
            sublist.append(word)            # add word to sublist
            if word == end:                 # the list met its end
                sublists.append(sublist)    # store it to sublists
                sublist = []                # empty it
        elif word == start:                 # met start word 
            sublist.append(word)            # add 1st word to sublist

    return sublists                         # return output
Related