# Collect subsequences of consecutive integers

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.

• Are all sequences sorted? – Ludisposed Feb 4 at 15:25
• @Ludisposed Yes – C. E. Feb 4 at 15:34

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


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

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()


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

• Thank you for pointing me to more-itertools! It has a function called consecutive_groups which does that which is needed. – C. E. Feb 5 at 9:11