Search code examples
pythonlistgroupingpython-itertools

Group consecutive integers and tolerate gaps of 1


In Python, given a list of sorted integers, I would to group them by consecutive values and tolerate gaps of 1.

For instance, given a list my_list:

In [66]: my_list
Out[66]: [0, 1, 2, 3, 5, 6, 10, 11, 15, 16, 18, 19, 20]

I would like the following output:

[[0, 1, 2, 3, 5, 6], [10, 11], [15, 16, 18, 19, 20]]

Now, if I didn't have to tolerate gaps of 1, I could apply the neat solution explained here:

import itertools
import operator
results = []
for k, g in itertools.groupby(enumerate(my_list), lambda (i,x):i-x):
        group = map(operator.itemgetter(1), g)
        results.append(group)

Is there a way to incorporate my extra requirement in the above solution? If not, what's the best way to tackle the problem?


Solution

  • When in doubt you can always write your own generator:

    def group_runs(li,tolerance=2):
        out = []
        last = li[0]
        for x in li:
            if x-last > tolerance:
                yield out
                out = []
            out.append(x)
            last = x
        yield out
    

    demo:

    list(group_runs(my_list))
    Out[48]: [[0, 1, 2, 3, 5, 6], [10, 11], [15, 16, 18, 19, 20]]