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!
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.