Identify groups of continuous numbers in a list

mikemaccana picture mikemaccana · Jan 28, 2010 · Viewed 56.5k times · Source

I'd like to identify groups of continuous numbers in a list, so that:

myfunc([2, 3, 4, 5, 12, 13, 14, 15, 16, 17, 20])

Returns:

[(2,5), (12,17), 20]

And was wondering what the best way to do this was (particularly if there's something inbuilt into Python).

Edit: Note I originally forgot to mention that individual numbers should be returned as individual numbers, not ranges.

Answer

Nadia Alramli picture Nadia Alramli · Jan 28, 2010

EDIT 2: To answer the OP new requirement

ranges = []
for key, group in groupby(enumerate(data), lambda (index, item): index - item):
    group = map(itemgetter(1), group)
    if len(group) > 1:
        ranges.append(xrange(group[0], group[-1]))
    else:
        ranges.append(group[0])

Output:

[xrange(2, 5), xrange(12, 17), 20]

You can replace xrange with range or any other custom class.


Python docs have a very neat recipe for this:

from operator import itemgetter
from itertools import groupby
data = [2, 3, 4, 5, 12, 13, 14, 15, 16, 17]
for k, g in groupby(enumerate(data), lambda (i,x):i-x):
    print map(itemgetter(1), g)

Output:

[2, 3, 4, 5]
[12, 13, 14, 15, 16, 17]

If you want to get the exact same output, you can do this:

ranges = []
for k, g in groupby(enumerate(data), lambda (i,x):i-x):
    group = map(itemgetter(1), g)
    ranges.append((group[0], group[-1]))

output:

[(2, 5), (12, 17)]

EDIT: The example is already explained in the documentation but maybe I should explain it more:

The key to the solution is differencing with a range so that consecutive numbers all appear in same group.

If the data was: [2, 3, 4, 5, 12, 13, 14, 15, 16, 17] Then groupby(enumerate(data), lambda (i,x):i-x) is equivalent of the following:

groupby(
    [(0, 2), (1, 3), (2, 4), (3, 5), (4, 12),
    (5, 13), (6, 14), (7, 15), (8, 16), (9, 17)],
    lambda (i,x):i-x
)

The lambda function subtracts the element index from the element value. So when you apply the lambda on each item. You'll get the following keys for groupby:

[-2, -2, -2, -2, -8, -8, -8, -8, -8, -8]

groupby groups elements by equal key value, so the first 4 elements will be grouped together and so forth.

I hope this makes it more readable.

python 3 version may be helpful for beginners

import the libraries required first

from itertools import groupby
from operator import itemgetter

ranges =[]

for k,g in groupby(enumerate(data),lambda x:x[0]-x[1]):
    group = (map(itemgetter(1),g))
    group = list(map(int,group))
    ranges.append((group[0],group[-1]))