문제

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.

라이센스 : CC-BY-SA ~와 함께 속성
제휴하지 않습니다 StackOverflow
scroll top