[Python] itertoolsモジュールを用いたイテレータの生成方法

Python モジュール 組込み関数

itertoolsモジュールは、イテレータの使い方としてよくある実装をライブラリにしたものです。例えば、イテラブルの要素の組み合わせパターンを全て取得したい場合などに便利です。このモジュールは多くの関数を実装していますが、本記事では主なものをいくつかピックアップしてまとめます。
# 2019/02/22 記事更新

確認した環境

  • OS: Ubuntu 16.04LTS
  • Python: ver3.7.0

組み合わせ(重複無し)の生成

itertools.combinations()は、イテラブル(文字列やリスト)の各要素の組み合わせパターンを全て出力します。
使い方は以下です。

itertools.combinations(iterable, r)

iterableから個の要素を抜き出した場合の組み合わせ(但し、重複は無し)をタプルで返します。返り値はイテレータオブジェクトとなります。

例えば、a、b、cのアルファベットが書かれたカードから2枚を取り出したときの組み合わせパターンを全て求めるには、以下のように書きます。

>>> import itertools
>>> comb = itertools.combinations('abc', 2)

# 返り値はイテレータオブジェクトです
>>> comb
<itertools.combinations object at 0x7fd667946188>

# 組み合わせはタプルで出力されます。以下は組み合わせパターンをリストで出力した場合 
>>> list(comb)
[('a', 'b'), ('a', 'c'), ('b', 'c')]

組み合わせ(重複有り)の生成

itertools.combinations_with_replacement()を使います。
使い方はcombinations()と同様です。

itertools.combinations_with_replacement(iterable, r)

例えば、a、b、cのアルファベットが書かれたカードから1枚ずつ取り出し、都度元に戻す操作を2回繰り返したとき(=重複有り)の組み合わせパターンを全て求めるには、以下のように書きます。

>>> import itertools
>>> comb = itertools.combinations_with_replacement('abc', 2)

# 返り値はイテレータオブジェクトです
>>> comb
<itertools.combinations_with_replacement object at 0x7fd6679461d8>

# 組み合わせはタプルで出力されます。以下は組み合わせパターンをリストで出力した場合
>>> list(comb)
[('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'b'), ('b', 'c'), ('c', 'c')]

順列の生成

itertools.permutations()を使うと、データの順列の全パターンを出力します。
使い方は以下です。

itertools.permutations(iterable, r=None)

iterableの要素からなる長さrの順列を返します。rが指定されない場合やNoneの場合はr = (iterableの長さ)となります。

例えば、a、b、cのアルファベットが書かれたカードから順番に1枚ずつ取り出た場合の組み合わせパターンを全て求めるには、以下のように書きます。

>>> import itertools
>>> per = itertools.permutations('abc', 2)

# 返り値はイテレータオブジェクトです
>>> per
<itertools.permutations object at 0x7fd66792c888>

# 組み合わせはタプルで出力されます。以下は組み合わせパターンをリストで出力した場合
>>> list(per)
[('a', 'b'), ('a', 'c'), ('b', 'a'), ('b', 'c'), ('c', 'a'), ('c', 'b')]

デカルト積(直積)

例えば、イテラブルAとイテラブルBの各要素の組み合わせを全て求めるには、for文を使って下記のように書けます。(※)

>>> A = ['a', 'b', 'c']
>>> B = ['0', '1']

>>> l = []
>>> for x in A:
...     for y in B:
...         l.append((x, y))
... 
>>> print(l)
[('a', '0'), ('a', '1'), ('b', '0'), ('b', '1'), ('c', '0'), ('c', '1')]

itertools.product()を使うと、上記をもっと簡単に書くことができます。
使い方は以下です。

itertools.product(*iterables, repeat=1)

これを使うと、上記の例は下記のように書けます。

>>> import itertools
>>> p = itertools.product(A, B)

# ちなみに、返り値はイテレータオブジェクトです
>>> p
<itertools.product object at 0x7fd667944e58>

# 組み合わせはタプルで出力されます。以下は組み合わせパターンをリストで出力した場合
>>> list(p)
[('a', '0'), ('a', '1'), ('b', '0'), ('b', '1'), ('c', '0'), ('c', '1')]

尚、キーワード引数repeatはその名の通り繰り返し回数を設定します。iterable自身の直積を求める際に使います。

>>> p = itertools.product(A, repeat=2)
>>> list(p)
[('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'a'), ('b', 'b'), ('b', 'c'), ('c', 'a'), ('c', 'b'), ('c', 'c')]

(※)itertools.product()はジェネレータ式の入れ子になった for ループとおよそ等価とのことですので、下記のほうが近いかもしれません。

# ジェネレータ式を使った例
>>> def prod(A, B):
...     for x in A:
...         for y in B:
...             yield (x, y)
... 
>>> list(prod(A, B))
[('a', '0'), ('a', '1'), ('b', '0'), ('b', '1'), ('c', '0'), ('c', '1')]

要素を抽出する

あるイテラブルについて、最初から5番目の要素のみを抽出するとか、1つおきに抽出するなどイテラブル対してスライス操作をする場合は、islice()メソッドを使います。

itertools.islice(iterable, stop)
itertools.islice(iterable, start, stop[, step])

start, stop, stepの考え方はスライス操作(参考記事)と同じです。

例を示します。

>>> s = 'abcdefg'

#第二引数以降が一つしかない場合はstopが設定される
>>> list(itertools.islice(s, 4))
['a', 'b', 'c', 'd']

#[2:4]のスライス操作と等価
>>> list(itertools.islice(s, 2, 4))
['c', 'd']

#startがNoneの場合は0、stopがNoneの場合はデータ末尾まで
>>> list(itertools.islice(s, None, None, 2))
['a', 'c', 'e', 'g']

一つのイテラブルから独立した複数のイテレータを生成

イテレータは通常、一度読み込むと再度読むことはでき無いため(参考記事)、同じイテレータに対して別の処理をすることは出来ません。(同じイテレータを再度生成できればその限りではありませんが。。)
tee()を使うと、元のイテレータから独立した複数のイテレータを生成できます。これにより、それぞれのイテレータに対して独立に処理をすることができます。
使い方は以下です。

itertools.tee(iterable, n=2)

どのように動作するか具体的に見てみます。

>>> A = 'abc'
>>> t = itertools.tee(A, 3)

# 返り値はイテレータオブジェクトです
>>> t
(<itertools._tee object at 0x7fd667945708>, <itertools._tee object at 0x7fd667945748>, <itertools._tee object at 0x7fd667945788>)

# それぞれの中身を確認すると、Aがコピーされていることが確認できます
>>> for _ in t:
...     list(_)
... 
['a', 'b', 'c']
['a', 'b', 'c']
['a', 'b', 'c']

連続した同じ要素をグループ分けする

あるイテラブルにおいて、ある条件に合致した要素毎にグループ分けをするためにはgroupby()メソッドを使います。まずは使い方から。

itertools.groupby(iterable, key=None)

第一引数iterable各要素に対し、keyパラメータで設定した条件に合致した要素でグループを作り、タプルのペア(key, グループ)のイテレータオブジェクトを作成して返します。key関数を省略した場合は要素そのものを返します。

尚、処理対象のイテラブルをあらかじめソートしないとグループを作成できません。
keyパラメータの値が変わるごとに休止や新しいグループを生成するためです。

以下に例を示します。

# 処理対象の文字列
>>> s = 'abcdabcdabcdabc'

# 文字でソート
>>> s_ = sorted(s)

# groupby()でグループ分け
>>> g = itertools.groupby(s_)

# 返り値はイテレータオブジェクト
>>> g
<itertools.groupby object at 0x7fd66792c8e0>

# 各グループもイテレータオブジェクト
>>> list(g)
[('a', <itertools._grouper object at 0x7fd667943940>), ('b', <itertools._grouper object at 0x7fd667943978>), ('c', <itertools._grouper object at 0x7fd6679436d8>), ('d', <itertools._grouper object at 0x7fd667943828>)]

# 実際にどうグループ分けされているかを確認
>>> def group(gr):
...     for k, g in gr:
...         yield (k, list(g))

>>> dict(group(g))
{'a': ['a', 'a', 'a', 'a'], 'b': ['b', 'b', 'b', 'b'], 'c': ['c', 'c', 'c', 'c'], 'd': ['d', 'd', 'd', 'd']}

まとめ

今回は、Pythonの組み込みライブラリであるitertoolsモジュールについて、主な関数をピックアップしてまとめました。要素の組み合わせ等を簡単に得ることができます。

Learn more...

書籍でもう少し詳しく学びたい場合はこちらもどうぞ。筆者もかなり参考にさせてもらっています!

シェアする
ひびきをフォローする
Hbk project