Search code examples
pythonlistsplitintegerindices

Python: split list of integers based on step between them


I have the following problem. Having a list of integers, I want to split it, into a list of lists, whenever the step between two elements of the original input list is not 1. For example: input = [0, 1, 3, 5, 6, 7], output = [[0, 1], [3], [5, 6, 7]]

I wrote the following function, but it's uggly as hell, and I was wondering if anyone of you guys would help me get a nicer solution. I tried to use itertools, but couldn't solve it.

Here's my solution:

def _get_parts(list_of_indices):
    lv = list_of_indices
    tuples = zip(lv[:-1], lv[1:])
    split_values = []
    for i in tuples:
        if i[1] - i[0] != 1:
            split_values.append(i[1])
    string = '/'.join([str(i) for i in lv])
    substrings = []
    for i in split_values:
        part = string.split(str(i))
        substrings.append(part[0])
        string = string.lstrip(part[0])
    substrings.append(string)
    result = []
    for i in substrings:
        i = i.rstrip('/')
        result.append([int(n) for n in i.split('/')])
    return result

Thanks a lot!


Solution

  • This works with any iterable

    >>> from itertools import groupby, count
    >>> inp = [0, 1, 3, 5, 6, 7]
    >>> [list(g) for k, g in groupby(inp, key=lambda i,j=count(): i-next(j))]
    [[0, 1], [3], [5, 6, 7]]