Collect subsequences of consecutive integers

Question

I've written a function to split a list into several lists with consecutive integers, e.g.:

extract_subsequences([1,2,3,7,8,9])

[[1, 2, 3], [7, 8, 9]]

extract_subsequences([1,2,3])

[[1, 2, 3]]

extract_subsequences([1,2,3, 10])

[[1, 2, 3], [10]]

extract_subsequences([1,2, 4,5, 8,9])

[[1, 2], [4, 5], [8, 9]]

This is the code that I came up with:

import numpy as np

def extract_subsequences(seq):
    def split_sequence(seq):
            for i in np.arange(1, len(seq)):
                if seq[i] != seq[i - 1] + 1:
                    break

            if i < len(seq) - 1:
                return seq[:i], seq[i:]
            else:
                if seq[-1] == seq[-2] + 1:
                    return seq, []
                else:
                    return seq[:-1], [seq[-1]]

    res = []
    last = seq
    while len(last) > 1:
        first, last = split_sequence(last)
        res.append(first)

    if len(last) == 1:
        res.append(last)

    return res

Can someone help me improve this code? It's kind of difficult to read with all the indices and special cases, and for the same reasons it was also quite difficult to write. I would very much appreciate a better way of thinking about this problem.

Answer 1

Using yield is often simpler than building up a list to return. Here it would reduce

    res = []
    last = seq
    while len(last) > 1:
        first, last = split_sequence(last)
        res.append(first)

    if len(last) == 1:
        res.append(last)

    return res

to

last = seq
while len(last) > 1:
    first, last = split_sequence(last)
    yield first

if len(last) == 1:
    yield last

at which point you could inline split_sequence.

---

split_sequence seems unnecessarily complicated to me. The entire extract_subsequences could be written with only one special case as

current = []
for val in seq:
    if current != [] and val != current[-1] + 1:
        yield current
        current = []
    current += [val]

# Test is only necessary because seq might be empty
if current != []:
    yield current

Answer 2

@Peter Taylor has a good idea, but there is more to be improved

• You should add some tests

• Python is often described as Batteries included

  This is a perfect example to use itertools.groupby()

---

from itertools import groupby
import doctest

def extract_seq(seq):
    """
    Splits sequence into consecutive lists
    
    args:
        seq (list): A sorted sequence

    >>> extract_seq([1,2, 4,5, 8,9])
    [[1, 2], [4, 5], [8, 9]]

    >>> extract_seq([1,2,3])
    [[1, 2, 3]]

    >>> extract_seq([1,2,3,7,8,9])
    [[1, 2, 3], [7, 8, 9]]

    >>> extract_seq([1,2,3,10])
    [[1, 2, 3], [10]]
    """
    return [
        [x for _, x in g]
        for k, g in groupby(
            enumerate(seq), 
            lambda i_x : i_x[0] - i_x[1]
        )
    ]

if __name__ == "__main__":
    doctest.testmod()

Answer 3

One more thing that no one has explicitly pointed out yet, because it is made irrelevant: there's no need to import numpy just to iterate over numpy.arange. Just use range instead: for i in range(1, len(seq)):. Or even better, use the itertools recipe pairwise (available with more-itertools):

for a, b in pairwise(seq):
    if b != a + 1:
        break