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

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

itertools --- 効率的なループ実行のためのイテレータ生成関数 — Python 3.10.0b2 ドキュメント

デカルト積とは各入力イテラブルから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.

問題文

NN 個の袋があります。
袋 ii には L_iLi 個のボールが入っていて、袋 ii の j(1\leq j\leq L_i)j(1≤jLi​) 番目のボールには正の整数 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)

コメント

タイトルとURLをコピーしました