iterable を分ける

この記事は Python Tips Advent Calendar 2012 １８日目の記事です。

Python の iterable オブジェクトから生成された iterator は一度しか走査できません。

>>> xs = iter([1,2,3])
>>> for x in xs: # xs の要素を全て消費
...     pass
>>> # もう消費したオブジェクトを戻せない
>>> xs.next() #doctest: +ELLIPSIS
Traceback (most recent call last):
  ...
StopIteration

しかし itertools.tee を使うことで、複数の独立した iterator に分けることができます。

>>> from itertools import tee
>>> xs = iter([1,2,3])
>>> ys,zs = tee(xs, 2) # 2 は分割数。省略した場合は 2 になる
>>> for y in ys:
...     print y
1
2
3
>>> for z in zs:
...     print z
1
2
3

ただし、tee は単に消費した要素をキューに詰めて再度使えるようにしているだけなので、この例のように、片方を全て消費してからもう片方を消費するという場合は、全ての要素をメモリに入れて列挙するのと同じだけのメモリを取ります。
この場合、ドキュメントにも書いてあるように、単に list に格納した上で使ったほうがいいでしょう。

tee は、例えば以下の様な、隣り合った要素を取り出すようなケースで有用です。

>>> from itertools import tee, izip, islice
>>> xs = iter(range(5))
>>> ys,zs = tee(xs)
>>> for y,z in izip(ys, islice(zs, 1, None)):
...     print y,z
0 1
1 2
2 3
3 4

これは ys と zs の進んでいる数が 1 個しか違わないので、1 要素分しかメモリを取りません。
たとえ xs の要素が無限リストであったとしても、これなら処理できます。

.