itertoolsモジュールをひと通り眺めてみた

Python 3 からは一部の関数がイテレータを返すよう変更になったということなので、この機会にPython3を触りつつitertoolsモジュールについてひととおり学んでみた。
主に参考にしたのは9.7. itertools — 効率的なループ実行のためのイテレータ生成関数 — Python 2.6ja2 documentation。そのため添えてあるメモは基本的にpython2.6が前提。2.7や3.2で追加されたり変更された点があればその旨を適宜書いてあります。

print関数について

今回はPython3.2であれこれしたのでprint関数を使ってます。Python2.6以降で print("", end=" ") みたいなのを使うには

from __future__ import print_function

とする必要があります。

バージョン毎のメソッドの実装状況

それぞれのバージョン毎で使えるメソッドが多少違うので表にしてみた。

3.2	2.7	2.6
accumulate	なし	なし
chain	chain	chain
combinations	combinations	combinations
combinations_with_replacement	combinations_with_replacement	なし
compress	compress	なし
count	count	count
cycle	cycle	cycle
dropwhile	dropwhile	dropwhile
groupby	groupby	groupby
組み込みのfilter関数	ifilter	ifilter
filterfalse	ifilterfalse	ifilterfalse
組み込みのmap関数	imap	imap
islice	islice	islice
組み込みのzip関数	izip	izip
zip_longest	izip_longest	izip_longest
permutations	permutations	permutations
product	product	product
repeat	repeat	repeat
starmap	starmap	starmap
takewhile	takewhile	takewhile
tee	tee	tee

というわけでまとめ

### 無限イテレータ

# itertools.count([n])
"""
引数nで始まる連続した整数を無限に返すイテレータを生成する。
python 2.7 からは
count(start, [step])
のようにステップ数を指定できるようになった。
"""
itr = itertools.count(5)
for i in itr:
    print(i, end=' ')
    if i == 10:
        break
# 5 6 7 8 9 10


# itertools.cycle(iterable)
"""
引数に与えられたイテレータの要素を次々と返す。
最後まで行くとStopIterationを起こさずに最初に戻って延々繰り返す。
これは最初にイテレータのコピーを保存してあるかららしい。
そのため与えるイテレータが大きい場合、メモリ消費量が増える。
"""
itr = itertools.cycle("ABC")
for i,j in enumerate(itr):
    print(j, end=' ')
    if i == 5:
        break
# A B C A B C


# itertools.repeat(object[, times])
"""
オブジェクトを指定した回数返すイテレータを生成する。
回数を指定しなければ無限に値を返すイテレータとなる。
"""
itr = itertools.repeat("YAH!", 3)
print([i for i in itr])
# ['YAH!', 'YAH!', 'YAH!']


### 一番短い入力シーケンスで止まるイテレータ


# itertools.chain(*iterables)
"""
引数に複数のイテレータを受け取り、それらがあたかもひとつのイテレータであるかのように
全ての要素を生成するイテレータを返す。
"""
itr = itertools.chain("ABC", range(5))
print(list(itr))
# ['A', 'B', 'C', 0, 1, 2, 3, 4]


# itertools.compress(data, selectors)
"""
dataとselectorsふたつのイテレータを受け取る。
セレクターが真になるデータを返すイテレータを生成する。
Python 2.7 で追加された
"""
itr = itertools.compress('ABCDEF', [1,0,1,0,0,1])
print(list(itr))
# ['A', 'C', 'F']


# itertools.dropwhile(predicate, iterable)
"""
predicateが真である限り無視し続け
その後はすべての要素を返すイテレータを返す。
"""
p = lambda x: x<5
l = [1,4,5,1,9,3,6]
itr = itertools.dropwhile(p, l)
print(list(itr))
# [5, 1, 9, 3, 6]


# iitertools.takewhile(predicate, iterable)
"""
predicateが真である限りiterableから要素を返すイテレータを返す。
"""
p = lambda x: x<5
l = [1,4,2,1,9,3,6]
itr = itertools.takewhile(p,l)
print(list(itr))
# [1, 4, 2, 1]


# itertools.groupby(iterable[, key])
"""
連続する要素をグループ化する。
値, <itertools._grouper object at 0x01232370>
尚、値が等しくても連続していなければグループ化されないので
必要に応じて先にソートしておく。
"""
for k, g in itertools.groupby(sorted("AAAABBCCCDDDEEEFFAAACC"), key=list):
    print(k, len([i for i in g]), end=', ')
# ['A'] 7, ['B'] 2, ['C'] 5, ['D'] 3, ['E'] 3, ['F'] 2,


# itertools.ifilter(predicate, iterable)
"""
predicateが真であるものだけを返す
"""
p = lambda x: x<5
l = [1,4,5,1,9,3,6]
# itr = itertools.ifilter(p,l)
itr = filter(p,l) # for Python3
print(list(itr))
# [1, 4, 1, 3]


# itertools.ifilterfalse(predicate, iterable)
"""
predicateが偽であるものだけを返す
"""
p = lambda x: x<5
l = [1,4,5,1,9,3,6]
# itr = itertools.ifilterfalse(p,l)
itr = itertools.filterfalse(p,l) # for Python3
print(list(itr))
# [5, 9, 6]


# itertools.islice(iterable[, start], stop[, step])
"""
イテレータを返すタイプのリストのスライス（list[start:stop:step])。
リスト同様、開始位置・終了位置・ステップを指定できる（stopは必須）
stopをNoneにすることで最後の要素まで返すが、渡したイテレータが
無限イテレータだと無限ループに陥るので注意。
"""
itr = itertools.islice(range(10), 3, None, 2)
print(list(itr))
# [3, 5, 7, 9]


# itertools.imap(function, *iterables)
"""
イテレータの要素を引数として関数を呼び出すイテレータを返す。
functionがNoneの場合は引数のタプルを返す。
Python3では2系で組み込みのmap関数がこのimapに置き換えられた。
"""
ls = [1,2,3,4,5]
# itr = itertools.imap(lambda x: x*2, ls)
itr = map(lambda x: x*2, ls) # for Python3
print(list(itr))
# [2, 4, 6, 8, 10]


# itertools.starmap(function, iterable)
"""
関数の引数が単一のイテレータにタプルとして格納されている場合は
imapの代わりにこっちを使う。
"""
tpls = [(2,3), (3,4), (6,8)]
itr = itertools.starmap(lambda x, y: x*y, tpls)
print(list(itr))
# [6, 12, 48]


# itertools.tee(iterable[, n=2])
"""
ひとつのイテレータから独立したn個（デフォルトでは2個）の独立したイテレータを返す。
どのサイトみても「Unixのteeコマンドに似てる」と書いてある。
"""
ls = [1,2,3,4,5]
tee1, tee2, = itertools.tee(ls)
print([i for i in tee1])   # [1, 2, 3, 4, 5]
print([i*2 for i in tee2]) # [2, 4, 6, 8, 10]

"""
なお、元のイテレータとteeで生成されたイテレータは共有関係にあるので
teeで新たなイテレータを生成した後はオリジナルのイテレータを使用してはいけない。
"""
ls = iter([1,2,3])
tee1, tee2, = itertools.tee(ls)
next(ls) #  1
print([i for i in tee1])   # [2, 3]
print([i*2 for i in tee2]) # [4, 6])


# itertools.izip(*iterables))
"""
各イテレータの要素をタプルに結合するイテレータを返す。
長さは引数に与えたイテレータの最も短いものに揃えられる。
溢れた要素は無視される。
Python3では2系で組み込みのzip関数がこのizipに置き換えられた。
"""
# itr = itertools.izip("ABCDE","xyz")
itr = zip("ABCDE","xyz") # for Python3
print(list(itr))
# [('A', 'x'), ('B', 'y'), ('C', 'z')]


# itertools.izip_longest(*iterables[, fillvalue])
"""
python3ではizip_longestからzip_longestへ名称変更。
izipとは逆に最も長いイテレータに長さを揃える。
足りない部分はfillvalueに値を指定すればそれで埋められる。
指定しなければNoneが入る。
"""
# itr = itertools.izip_longest("ABCDE","xyz", fillvalue='-')
itr = itertools.zip_longest("ABCDE","xyz", fillvalue='-')
print(list(itr))
# [('A', 'x'), ('B', 'y'), ('C', 'z'), ('D', '-'), ('E', '-')]


### 組み合わせジェネレータ


# itertools.product(*iterables[, repeat])
"""
デカルト積というやつらしい。

list(itertools.product("ABC", range(3)))

↑↓ ネストしたループと等価。

[(j,i) for i in range(3) for j in "ABC"]
"""
product_test = itertools.product("AB", "12", repeat=2)
print(list(itr))
"""
[('A', '1', 'A', '1'), ('A', '1', 'A', '2'), ('A', '1', 'B', '1'),
('A', '1', 'B', '2'), ('A', '2', 'A', '1'), ('A', '2', 'A', '2'),
('A', '2', 'B', '1'), ('A', '2', 'B', '2'), ('B', '1', 'A', '1'),
('B', '1', 'A', '2'),('B', '1', 'B', '1'), ('B', '1', 'B', '2'),
('B', '2', 'A', '1'), ('B', '2', 'A', '2'), ('B', '2', 'B', '1'),
('B', '2', 'B', '2')]
"""


# itertools.permutations(iterable[, r])
"""
順列を求める。rは一度に取り出す要素数を指定する。
rを指定しなかったりNoneを与えると、rはiterableの長さとなる。
それぞれの順列はタプルで格納される。
"""
print(list(itertools.permutations([1,2,3], 2)))
# [(1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)]


# itertools.combinations(iterable, r)
"""
こちらは組み合わせ。重複は持たない。
"""
print(list(itertools.combinations([1,2,3,4], 3)))
# [(1, 2, 3), (1, 2, 4), (1, 3, 4), (2, 3, 4)]


# itertools.combinations_with_replacement(iterable, r)
"""
重複を持った組み合わせはこっちを使う。
Python 2.7から追加された
"""
print(list(itertools.combinations_with_replacement([1,2,3,4], 3)))

"""
[(1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 1, 4), (1, 2, 2),
(1, 2, 3), (1, 2, 4), (1, 3, 3), (1, 3, 4), (1, 4, 4),
(2, 2, 2), (2, 2, 3), (2, 2, 4), (2, 3, 3),(2, 3, 4),
(2, 4, 4), (3, 3, 3), (3, 3, 4), (3, 4, 4), (4, 4, 4)]
"""


# itertools.accumulate(iterable)
"""
Numpyのaccumulateがモデルらしい（Numpy触ったことない）。
要素を取り出した時点での累計が返る。

Python 3.2から追加された
"""
print(list(itertools.accumulate([2,9,57,10,32])))
# [2, 11, 68, 78, 110]

もっとカッコイイ例が書けるようになりたい。