[python]itertoolsを使って多重ループをする方法

python_itertools.ipynb

{"cells":[{"cell_type":"code","execution_count":1,"metadata":{},"outputs":[],"source":["from itertools import product, permutations, combinations, combinations_with_replacement\n","import itertools\n","from math import factorial\n","from scipy.special import perm, comb\n",""]},{"cell_type":"markdown","metadata":{},"source":[" 長さNのR重ループはforをR回ネストすることで実現できるが、`itertools.product`で同様の処理ができる。\n"," 例えばN=5, R=3の場合以下のようになる。計算量は$O(n^r)$となる。"]},{"cell_type":"code","execution_count":2,"metadata":{},"outputs":[{"output_type":"execute_result","data":{"text/plain":["(125, 750)"]},"metadata":{},"execution_count":2}],"source":["# 長さNの配列の3重ループ\n","N = 5\n","cnt = 0\n","ans = 0\n","for i in range(N):\n","    for j in range(N):\n","        for k in range(N):\n","            cnt += 1\n","            ans += i + j + k\n","cnt, ans\n",""]},{"cell_type":"code","execution_count":3,"metadata":{},"outputs":[{"output_type":"execute_result","data":{"text/plain":["(125, 750)"]},"metadata":{},"execution_count":3}],"source":["cnt = 0\n","ans = 0\n","for i, j, k in product(range(N), repeat=3):\n","    cnt += 1\n","    ans += i + j + k\n","cnt, ans\n",""]},{"cell_type":"markdown","metadata":{},"source":[" この場合、i,j,kは重複する。重複させたくない場合は以下のようにする。"]},{"cell_type":"code","execution_count":4,"metadata":{},"outputs":[{"output_type":"stream","name":"stdout","text":["2 1 0\n3 1 0\n3 2 0\n3 2 1\n4 1 0\n4 2 0\n4 2 1\n4 3 0\n4 3 1\n4 3 2\n"]},{"output_type":"execute_result","data":{"text/plain":["(10, 60)"]},"metadata":{},"execution_count":4}],"source":["N = 5\n","cnt = 0\n","ans = 0\n","for i in range(N):\n","    for j in range(i):  # Nをiに変更、iが0なら処理は行われない\n","        for k in range(j):  # Nをjに変更、jが0なら処理は行われない\n","            print(i, j, k)\n","            cnt += 1\n","            ans += i + j + k\n","cnt, ans"]},{"cell_type":"code","execution_count":5,"metadata":{},"outputs":[{"output_type":"stream","name":"stdout","text":["0 1 2\n0 1 3\n0 1 4\n0 2 3\n0 2 4\n0 3 4\n1 2 3\n1 2 4\n1 3 4\n2 3 4\n"]},{"output_type":"execute_result","data":{"text/plain":["(10, 60)"]},"metadata":{},"execution_count":5}],"source":["cnt = 0\n","ans = 0\n","for i, j, k in combinations(range(N), 3):\n","    print(i, j, k)\n","    cnt += 1\n","    ans += i + j + k\n","cnt, ans"]},{"cell_type":"markdown","metadata":{},"source":[" その総数(=計算量)の求め方は\n"," $$\n"," \\Large\\frac{n!}{r! \\times (n - r)!}\n"," $$\n"," pythonではmath.factorialやscipy.special.combが使える。"]},{"cell_type":"code","execution_count":6,"metadata":{},"outputs":[{"output_type":"execute_result","data":{"text/plain":["(10, 10.0)"]},"metadata":{},"execution_count":6}],"source":["a = factorial(5) // (factorial(5 - 3) * factorial(3))\n","b = comb(5, 3)\n","a, b\n",""]},{"cell_type":"markdown","metadata":{},"source":[" itertoolsにはproduct, permutations, combinations, combinations_with_replacementの各関数が有り以下の通りに作動する"]},{"cell_type":"code","execution_count":7,"metadata":{},"outputs":[{"output_type":"stream","name":"stdout","text":["product:  [(1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 2, 1), (1, 2, 2), (1, 2, 3), (1, 3, 1), (1, 3, 2), (1, 3, 3), (2, 1, 1), (2, 1, 2), (2, 1, 3), (2, 2, 1), (2, 2, 2), (2, 2, 3), (2, 3, 1), (2, 3, 2), (2, 3, 3), (3, 1, 1), (3, 1, 2), (3, 1, 3), (3, 2, 1), (3, 2, 2), (3, 2, 3), (3, 3, 1), (3, 3, 2), (3, 3, 3)]\nlen:  27\n"]}],"source":["li = list(product([1, 2, 3], repeat=3))\n","print(\"product: \", li)\n","print(\"len: \", len(li))\n",""]},{"cell_type":"code","execution_count":8,"metadata":{},"outputs":[{"output_type":"stream","name":"stdout","text":["permutations:  [(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)]\nlen:  6\n"]}],"source":["li = list(permutations([1, 2, 3]))\n","print(\"permutations: \", li)\n","print(\"len: \", len(li))\n","\n",""]},{"cell_type":"code","execution_count":9,"metadata":{},"outputs":[{"output_type":"stream","name":"stdout","text":["combinations:  [(1, 2, 3)]\n","len:  1\n"]}],"source":["li = list(combinations([1, 2, 3], 3))\n","print(\"combinations: \", li)\n","print(\"len: \", len(li))\n",""]},{"cell_type":"code","execution_count":10,"metadata":{},"outputs":[{"output_type":"stream","name":"stdout","text":["combinations_with_replacement:  [(1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 2, 2), (1, 2, 3), (1, 3, 3), (2, 2, 2), (2, 2, 3), (2, 3, 3), (3, 3, 3)]\nlen:  10\n"]}],"source":["li = list(combinations_with_replacement([1, 2, 3], 3))\n","print(\"combinations_with_replacement: \", li)\n","print(\"len: \", len(li))\n","\n",""]}],"nbformat":4,"nbformat_minor":2,"metadata":{"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":3},"orig_nbformat":2}}