itertools.groupbyを使ってみる

Rubyの先頭のアルファベット別に要素をグループ分けするサンプルを書き換えてみる。
が、ちょっとコケた。
groupbyの挙動を見てみよう。

from collections import defaultdict
from itertools import groupby
animals = ["cat", "bat", "bear", "camel", "alpaca"]
d = defaultdict(list)
for k, g in groupby(animals, key=lambda x: x[0]):
    d[k].append(list(g))
    print list(g)

""" print list(g)の出力結果
['cat']
['bat', 'bear'] <--これ
['camel']
['alpaca']
"""
print d.items()
# => [('a', [['alpaca']]), ('c', [['cat'], ['camel']]), ('b', [['bat', 'bear']])]

list中でkey要素が連続する場合、gは複数要素を持っているらしい。
この点に注意して処理を書く。下はStopIterationを捕捉するまでg.next()を繰り返している。

from collections import defaultdict
from itertools import groupby
animals = ["cat", "bat", "bear", "camel", "alpaca"]
d = defaultdict(list)
for k, g in groupby(animals, key=lambda x: x[0]):
    try:
        while g:
            d[k].append(g.next())
    except StopIteration:
        pass

print d.items()
# => [('a', ['alpaca']), ('c', ['cat', 'camel']), ('b', ['bat', 'bear'])]

しかしいちいち例外処理を回していては遅いはずなので前処理でsortedするべき。

from collections import defaultdict
from itertools import groupby
animals = ["cat", "bat", "bear", "camel", "alpaca"]
d = defaultdict(list)
animals = sorted(animals, key=lambda x: x[0])
for k, g in groupby(animals, key=lambda x: x[0]):
    d[k].append(list(g))
    # リストの入れ子をはずす
    d[k] = d[k][0]

print d.items()
# => [('a', ['alpaca']), ('c', ['cat', 'camel']), ('b', ['bat', 'bear'])]

d[k] = d[k][0]はちょっと汚いが問題なく動く。