4
\$\begingroup\$

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.

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

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
\$\endgroup\$
9
\$\begingroup\$

@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()
\$\endgroup\$
4
\$\begingroup\$

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
\$\endgroup\$
  • 2
    \$\begingroup\$ Thank you for pointing me to more-itertools! It has a function called consecutive_groups which does that which is needed. \$\endgroup\$ – C. E. Feb 5 at 9:11

Your Answer

By clicking "Post Your Answer", you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.