compact-number-list (2)

• 6.5.2 Exampla (http://docs.python.org/lib/itertools-example.html)

Pythonのマニュアルにめちゃくちゃうまいやり方が載っていた。

>>> data = [ 1,  4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
...     print map(operator.itemgetter(1), g)
... 
[1]
[4, 5, 6]
[10]
[15, 16, 17, 18]
[22]
[25, 26, 27, 28]

以下は、上記のやり方を利用したバージョン。

from operator import itemgetter
from itertools import groupby

def compact_number_list(seq):
    for k, g in groupby(enumerate(seq), lambda (i,x):i-x):
        L = map(itemgetter(1), g)
        yield (L[0], [L[0], L[-1]])[len(L) != 1]

compact_number_list.__doc__ = \
"""
>>> data = [1, 3, 4, 5, 6, 12, 13, 15]
>>> list(compact_number_list(data))
[1, [3, 6], [12, 13], 15]
"""

if __name__ == '__main__':
    import doctest
    doctest.testmod()