【Python itertools】総当たりで組み合わせを求める

やりたいこと

以下のように与えられた配列の各要素を超えないような数のすべての組み合わせを求めたいとします。

Args: [3, 1, 2]
Returns: [
    [1, 1, 1],
    [1, 1, 2],
    [2, 1, 1],
    [2, 1, 2],
    [3, 1, 1],
    [3, 1, 2]
]

実装

from itertools import product
from string import ascii_lowercase

def get_combination(input_list):
    return [
        list(map(lambda x:ascii_lowercase.index(x) + 1, tup))
        for tup in list(product(*tuple([ascii_lowercase[:num] for num in input_list])))
    ]

if __name__ == '__main__':
    combination = get_combination([3, 1, 2])
    print(combination)

出力は以下のようになります。

[[1, 1, 1], [1, 1, 2], [2, 1, 1], [2, 1, 2], [3, 1, 1], [3, 1, 2]]

解説

string.ascii_lowercase で小文字アルファベットの文字列を取得します。これをインデックスで切り出します。例えば以下のようになります。

>>> from string import ascii_lowercase
>>> ascii_lowercase
'abcdefghijklmnopqrstuvwxyz'
>>> ascii_lowercase[:2]
'ab'

itertools.product で総当たりの組み合わせを求めます。例えば以下のようになります。

>>> from itertools import product
>>> list(product('ab', 'ab'))
[('a', 'a'), ('a', 'b'), ('b', 'a'), ('b', 'b')]

('a', 'b') を [1, 2] に変換する方法は以下のようになります。

>>> list(map(lambda x: ascii_lowercase.index(x) + 1, ('a', 'b')))
[1, 2]

参考資料