【python】itertools.productでデカルト積を求める

itertools.productで入力イテラブルのデカルト積を求めることができます。

itertools --- 効率的なループ実行のためのイテレータ生成関数

このモジュールは イテレータ を構築する部品を実装しています。プログラム言語 APL, Haskell, SML からアイデアを得ていますが、 Python に適した形に修正されています。 このモジュールは、高速でメモリ効率に優れ、単独でも組合せても使用することのできるツールを標準化したものです。同時に、このツール群は...

docs.python.org

デカルト積とは各入力イテラブルから1つずつ要素を取り出したすべての組み合わせのことです。

A = [1, 2, 3]
B = [4, 5, 6]
C = [7, 8, 9]

このデカルト積は

(1, 4, 7), (1, 4, 8), (1, 4, 9), (1, 5, 7), (1, 5, 8), (1, 5, 9), (1, 6, 7), (1, 6, 8), (1, 6, 9)
(2, 4, 7), (2, 4, 8), (2, 4, 9), (2, 5, 7), (2, 5, 8), (2, 5, 9), (2, 6, 7), (2, 6, 8), (2, 6, 9)
(3, 4, 7), (3, 4, 8), (3, 4, 9), (3, 5, 7), (3, 5, 8), (3, 5, 9), (3, 6, 7), (3, 6, 8), (3, 6, 9)

こんな感じです。

コードと実行結果は

from itertools import product

A = [1, 2, 3]
B = [4, 5, 6]
C = [7, 8, 9]

for p in product(A, B, C):
    print(p)

実行結果

(1, 4, 7)
(1, 4, 8)
(1, 4, 9)
(1, 5, 7)
(1, 5, 8)
(1, 5, 9)
(1, 6, 7)
(1, 6, 8)
(1, 6, 9)
(2, 4, 7)
(2, 4, 8)
(2, 4, 9)
(2, 5, 7)
(2, 5, 8)
(2, 5, 9)
(2, 6, 7)
(2, 6, 8)
(2, 6, 9)
(3, 4, 7)
(3, 4, 8)
(3, 4, 9)
(3, 5, 7)
(3, 5, 8)
(3, 5, 9)
(3, 6, 7)
(3, 6, 8)
(3, 6, 9)

atcoderのABC233、C – Productにてこれが使えるようなので試します。

C - Product

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

atcoder.jp

問題文

NN 個の袋があります。
袋 ii には L_iLi​ 個のボールが入っていて、袋 ii の j(1\leq j\leq L_i)j(1≤j≤Li​) 番目のボールには正の整数 a_{i,j}ai,j​ が書かれています。

それぞれの袋から 11 つずつボールを取り出します。
取り出したボールに書かれた数の総積が XX になるような取り出し方は何通りありますか？

ただし、書かれた数が同じであっても全てのボールは区別します。

from itertools import product

n, x = map(int, input().split())
a = [list(map(int, input().split()))[1:] for i in range(n)]

ans = 0

print(a)
print(*a)

for p in product(*a):
    print(p)
    pr = 1
    for pi in p:
        pr *= pi
    if pr == x:
        ans += 1

print(ans)