Question

I have this piece of code to find consecutive numbers in a list:

from itertools import groupby
from operator import itemgetter

a = [1,2,3,5,55,56]

def consc(b):
  for k, g in groupby(enumerate(b), lambda (i,x):i-x):
     print map(itemgetter(1), g)  

consc(a)

output:

[1, 2, 3]
[5]
[55, 56]

However, I'd like to be able to seek for other deltas as well (1 to 10), for example a difference of 2 will yield the following output from the same list:

[1]
[2]
[3,5]
[55]
[56]

Thanks!

Was it helpful?

Solution

It's actually a pretty simple modification:

from itertools import groupby, count
from operator import itemgetter

a = [1,2,3,5,55,56]

def consc(b, step):
  for k, g in groupby(zip(count(step=step), b), lambda (i, x): i-x):
     print map(itemgetter(1), g)

consc(a, 2)

Which gives:

[1]
[2]
[3, 5]
[55]
[56]

Instead of using enumerate(), we use zip() and count() with a step of the desired value, which gives the wanted result.

Cleaned up a little:

from itertools import groupby, count
from operator import itemgetter

def _sub(item):
    a, b = item
    return a - b

def consecutive(iterable, step):
    for _, g in groupby(zip(count(step=step), iterable), _sub):
        yield map(itemgetter(1), g)

a = [1, 2, 3, 5, 55, 56]

print(list(consecutive(a, 2)))

It makes sense to have a generator here, and use more descriptive names. Using an actual function avoids re-declaring it every time the function is used, as with lambda. This also works in Python 3.x by avoiding using argument unpacking, which has been removed from the language.

Licensed under: CC-BY-SA with attribution
Not affiliated with StackOverflow
scroll top