note.nkmk.me

Pythonで複数のリストの直積（デカルト積）を生成するitertools.product

Modified: | Tags: Python, リスト

Pythonで複数のリストの直積（デカルト積）を生成するにはitertools.product()を使う。

目次

単独のリストの順列・組み合わせを生成するには、同じくitertoolsモジュールのitertools.permutations(), itertools.combinations()を使う。

本記事のサンプルコードでは、以下のようにitertoolsモジュールをインポートする。標準ライブラリなので追加のインストールは不要。pprintは結果を見やすくするために使用している。

import itertools
import pprint
source: itertools_product.py

直積（デカルト積）とは

直積（デカルト積）は、複数の集合から要素を一つずつ取り出した組み合わせの集合。

例えば2つのリストがあったとき、すべてのペアの組み合わせのリストが直積。以降に具体例を示す。

itertools.product()の基本的な使い方

返り値はイテレータ

引数に2つのリストを指定すると、itertools.product型のオブジェクトが返される。itertools.product型はイテレータなので、そのままprint()しても中身は出力されない。

l1 = [1, 2, 3]
l2 = ['A', 'B']

p = itertools.product(l1, l2)
print(p)
# <itertools.product object at 0x105e6e2c0>

print(type(p))
# <class 'itertools.product'>
source: itertools_product.py

forループで各リストの要素の組み合わせをタプルで取得できる。最後まで辿り着いたイテレータを再度forループでまわしても何も出力されないので注意。

for v in p:
    print(v)
# (1, 'A')
# (1, 'B')
# (2, 'A')
# (2, 'B')
# (3, 'A')
# (3, 'B')

for v in p:
    print(v)
source: itertools_product.py

タプルではなくそれぞれの要素を別々に取得することも可能。

for v1, v2 in itertools.product(l1, l2):
    print(v1, v2)
# 1 A
# 1 B
# 2 A
# 2 B
# 3 A
# 3 B
source: itertools_product.py

各リストをネストしたforループ（多重ループ）と同じ結果が得られる。

for v1 in l1:
    for v2 in l2:
        print(v1, v2)
# 1 A
# 1 B
# 2 A
# 2 B
# 3 A
# 3 B
source: itertools_product.py

結果をリスト化: list()

list()で結果をリストに変換することもできる。タプルを要素とするリストとなる。

l1 = [1, 2, 3]
l2 = ['A', 'B']

l_p = list(itertools.product(l1, l2))
print(l_p)
# [(1, 'A'), (1, 'B'), (2, 'A'), (2, 'B'), (3, 'A'), (3, 'B')]

print(type(l_p))
# <class 'list'>

print(type(l_p[0]))
# <class 'tuple'>
source: itertools_product.py

複数のイテラブルオブジェクトを指定可能

ここまでの説明では引数を便宜上「リスト」としてきたが、イテラブルオブジェクトであればタプルでも辞書でもrangeオブジェクトでもなんでもいい。

また、引数にはイテラブルオブジェクトを2個以上指定できる。

t = ('A', 'B')
d = {'key1': 'value1', 'key2': 'value2'}
r = range(2)

pprint.pprint(list(itertools.product(t, d, r)))
# [('A', 'key1', 0),
#  ('A', 'key1', 1),
#  ('A', 'key2', 0),
#  ('A', 'key2', 1),
#  ('B', 'key1', 0),
#  ('B', 'key1', 1),
#  ('B', 'key2', 0),
#  ('B', 'key2', 1)]
source: itertools_product.py

結果から分かるように、辞書オブジェクトをそのままイテレートするとキーkeyが使われる。値valueを使いたい場合はvalues()メソッドを使う。以下の記事を参照。

range()についての詳細は以下の記事を参照。

同じリストを繰り返し使用: 引数repeat

キーワード引数repeatに繰り返しの回数を指定すると、イテラブルオブジェクトを繰り返し使用して直積を生成する。

l1 = ['A', 'B']

pprint.pprint(list(itertools.product(l1, repeat=3)))
# [('A', 'A', 'A'),
#  ('A', 'A', 'B'),
#  ('A', 'B', 'A'),
#  ('A', 'B', 'B'),
#  ('B', 'A', 'A'),
#  ('B', 'A', 'B'),
#  ('B', 'B', 'A'),
#  ('B', 'B', 'B')]
source: itertools_product.py

repeatを指定しない次の形と同じ（デフォルトはrepeat=1）。

pprint.pprint(list(itertools.product(l1, l1, l1)))
# [('A', 'A', 'A'),
#  ('A', 'A', 'B'),
#  ('A', 'B', 'A'),
#  ('A', 'B', 'B'),
#  ('B', 'A', 'A'),
#  ('B', 'A', 'B'),
#  ('B', 'B', 'A'),
#  ('B', 'B', 'B')]
source: itertools_product.py

イテラブルオブジェクトを複数指定した場合は次のようになる。

l1 = [1, 2]
l2 = ['A', 'B']

pprint.pprint(list(itertools.product(l1, l2, repeat=2)))
# [(1, 'A', 1, 'A'),
#  (1, 'A', 1, 'B'),
#  (1, 'A', 2, 'A'),
#  (1, 'A', 2, 'B'),
#  (1, 'B', 1, 'A'),
#  (1, 'B', 1, 'B'),
#  (1, 'B', 2, 'A'),
#  (1, 'B', 2, 'B'),
#  (2, 'A', 1, 'A'),
#  (2, 'A', 1, 'B'),
#  (2, 'A', 2, 'A'),
#  (2, 'A', 2, 'B'),
#  (2, 'B', 1, 'A'),
#  (2, 'B', 1, 'B'),
#  (2, 'B', 2, 'A'),
#  (2, 'B', 2, 'B')]
source: itertools_product.py

次の形と同じ。l1, l1, l2, l2ではなくl1, l2, l1, l2となるので注意。

pprint.pprint(list(itertools.product(l1, l2, l1, l2)))
# [(1, 'A', 1, 'A'),
#  (1, 'A', 1, 'B'),
#  (1, 'A', 2, 'A'),
#  (1, 'A', 2, 'B'),
#  (1, 'B', 1, 'A'),
#  (1, 'B', 1, 'B'),
#  (1, 'B', 2, 'A'),
#  (1, 'B', 2, 'B'),
#  (2, 'A', 1, 'A'),
#  (2, 'A', 1, 'B'),
#  (2, 'A', 2, 'A'),
#  (2, 'A', 2, 'B'),
#  (2, 'B', 1, 'A'),
#  (2, 'B', 1, 'B'),
#  (2, 'B', 2, 'A'),
#  (2, 'B', 2, 'B')]
source: itertools_product.py

多重ループ（ネストしたループ）との速度比較

上述のように、itertools.product()で生成されるイテレータをforループでまわすと多重forループ（ネストしたforループ）と同じ結果となる。

l1 = [1, 2, 3]
l2 = ['A', 'B']

for v1, v2 in itertools.product(l1, l2):
    print(v1, v2)
# 1 A
# 1 B
# 2 A
# 2 B
# 3 A
# 3 B

for v1 in l1:
    for v2 in l2:
        print(v1, v2)
# 1 A
# 1 B
# 2 A
# 2 B
# 3 A
# 3 B
source: itertools_product.py

以降に示すように、実はitertools.product()のほうが多重forループよりも遅い。

イテラブルオブジェクトの要素数、ネストするforループの数によって結果が異なる可能性があるが、例えば以下のStack OverflowのQ&Aでもitertools.product()のほうが遅いという結果が回答されている。

以下、Jupyter Notebookのマジックコマンド%%timeitで実行時間を計測した結果を示す。Pythonコードとして実行しても計測できないので注意。

要素数1000個、二重ループ

要素数1000個、二重ループの例。pass文でスルーしており、ループを回しているだけ。

itertools.product()の結果はアンパックして受け取ったほうが速い。

A = range(1000)

%%timeit
for x in itertools.product(A, A):
    pass
# 30.8 ms ± 910 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
for a1, a2 in itertools.product(A, A):
    pass
# 22.8 ms ± 293 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
source: itertools_product_timeit.py

ネストしたforループはアンパックした場合のitertools.product()とほぼ同等（若干速い）。

%%timeit
for a1 in A:
    for a2 in A:
        pass
# 22.6 ms ± 345 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
source: itertools_product_timeit.py

リスト内包表記のジェネレータ版であるジェネレータ式を使う場合はアンパックしないほうが速いが、いずれにせよitertools.product()、ネストしたforループよりも遅い。

%%timeit
for x in ((a1, a2) for a1 in A for a2 in A):
    pass
# 82.2 ms ± 467 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
for a1, a2 in ((a1, a2) for a1 in A for a2 in A):
    pass
# 91.4 ms ± 276 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
source: itertools_product_timeit.py

具体的な例として各組み合わせの積の合計を算出する。この場合もitertools.product()よりもネストしたforループのほうが速い。

%%timeit
v = 0
for a1, a2 in itertools.product(A, A):
    v += a1 * a2
# 98.8 ms ± 579 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
v = 0
for a1 in A:
    for a2 in A:
        v += a1 * a2
# 95.7 ms ± 4.05 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
source: itertools_product_timeit.py

この例ではジェネレータ式をsum()に渡す方法のほうがわずかに速い。

%%timeit
v = sum(a1 * a2 for a1, a2 in itertools.product(A, A))
# 94 ms ± 2.36 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
v = sum(a1 * a2 for a1 in A for a2 in A)
# 92.7 ms ± 4.83 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
source: itertools_product_timeit.py

要素数100個、三重ループ

要素数100個、三重ループの例。

この場合もネストしたforループが一番速い。

B = range(100)

%%timeit
for x in itertools.product(B, B, B):
    pass
# 31.6 ms ± 725 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
for b1, b2, b3 in itertools.product(B, B, B):
    pass
# 26.2 ms ± 490 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
for b1 in B:
    for b2 in B:
        for b3 in B:
            pass
# 12.9 ms ± 176 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%%timeit
for x in ((b1, b2, b3) for b1 in B for b2 in B for b3 in B):
    pass
# 80.9 ms ± 1.27 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
for b1, b2, b3 in ((b1, b2, b3) for b1 in B for b2 in B for b3 in B):
    pass
# 93.8 ms ± 3.22 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
source: itertools_product_timeit.py

以上のように、要素数が1000個の二重ループや100個の三重ループでの差は数十ms程度。速度が重要な場合は多重ループ、特に速度を重視しない場合であればitertools.product()を使ってコードの見通しをよくする、というように状況に応じて使い分ければよいだろう。

当然ながら条件によって結果が異なる可能性があるので、どちらを採用するかは実際の条件で試して判断してほしい。

関連カテゴリー

関連記事