二次元リストのそれぞれの要素の組み合わせを抽出

Question

二次元リストのそれぞれの要素数が異なり、その要素を組み合わせた新たなリストを作りたいです。
具体的なコードを以下に示します。

A = [[1], [2, 3, 4], [5]]

という二次元リストがあるときに新しく

lst = [[1,2,5], [1,3,5], [1,4,5]]

というリストを作成したいです。
それぞれの要素数は任意な値です。

二次元リストではなく、複数のリストならばitertools.product()で生成可能だと思うのですが、二次元リストの各要素のリストが組み合わせの対象の場合、どのように対応すれば良いのでしょうか？

よろしくお願いします。

Answer 1

from itertools import product

A = [[1], [2, 3, 4], [5]]
lst = [*map(list, product(*A))]
print(lst) # [[1, 2, 5], [1, 3, 5], [1, 4, 5]]

Comments

[hirochem]

ありがとうございます！そういう書き方もできるんですね。
知りませんでした

Answer 2

二次元リストも複数のリストも本質的には（展開（アンパック）すれば）同じで、productで生成できます。

from itertools import product

A = [[1], [2, 3, 4], [5]]
lst = [list(v) for v in product(*A)]
print(lst) # [[1, 2, 5], [1, 3, 5], [1, 4, 5]]

Comments

[hirochem]

アンパックという言葉を初めて知りました。勉強になります。
一つ質問なのですが、「import itertools」でproductを使う場合と、「from itertools import product」でproductを使う場合の違いはどういったものでしょうか？

[can110]

記法の違い程度でとくに違いはありません。

[hirochem]

早速ありがとうございます！

助かりました。

Answer 3

import itertools
A = [[1], [2, 3, 4], [5]]

itr = itertools.product(*A)
lst = [list(i) for i in itr]

でどうでしょうか？

Comments

[hirochem]

数時間悩み続けたので非常に助かります！
まだまだ自分の力不足です。
map関数を使わなくても内胞表記すればタプルではなく、リストとして返せるんですね。
ありがとうございます！