Pythonで多重forループを避ける簡単な手法

はじめに

プログラミングにおいて多重forループ（ネスト） はもっとも避けたい処理のうちの1つだと思います。一か月ほど前に「なぜ我々は頑なにforを避けるのか」という記事がバズっており、2020/11/3現在で1000以上のLGTMが付いていることからも、多くの人はforループをあまり使いたくないのではないでしょうか。
Pythonには簡単に多重forループを避けることができる手法がいくつか存在するので、自分がよく使うものをまとめてみました。他にもあれば、ぜひコメントで教えてください。

そもそも、なぜ多重forループを避けたいのか

人それぞれいろんな理由があると思いますが（上記の記事にも色々詳しく考察されています）、個人的には以下の2点です。

1. 読みづらい

多重forループの例題として、3桁それぞれで0～9のfor文を回すことで0～999を全探索して、500以上かつゾロ目である数字を探すプログラムを書いてみます。（適当）
3重forループだけでもウッとなりますね。

悪い例

result = []
for i in range(10):
    for j in range(10):
        for k in range(10):
            summed = i*100+j*10+k
            if i == j == k and summed >= 500:
                result.append(summed)
print(result)

# 実行結果 -> [555, 666, 777, 888, 999]

処理が単純なのでまだマシですが、もっと複雑になるとほんとに読みづらくなりますね。

2. ループをまとめて抜ける処理が、非常に読みづらい

ある条件を満たしたら、一番内側のループだけでなくて全てのループを抜け出したい時ってよくありますよね。例題として、上記1.の処理で、ゾロ目を1つでも見つけたら処理を終了するようなプログラムを書きます。Googleで「python 多重ループ 抜け方」と検索してみると、flagを使う使わないといった差はあれど、おおよそ以下のような書き方をしています。
※Pythonで多重forループを抜けるシンプルな方法は、外部ライブラリのgotoくらいしかないため

悪い例breakあり

result = []
for i in range(10):
    for j in range(10):
        for k in range(10):
            summed = i*100 + j*10 + k
            if i == j == k and summed >= 500:
                result.append(summed)
                break
        else:
            continue
        break
    else:
        continue
    break

print(result)

# 実行結果 -> [555]

これはかなり読みづらいですね。どことどこのインデントが同じかがとても分かりづらいです。
より複雑な処理を書こうとするならば、だれにも理解できないコードが完成することでしょう。

手法A. itertoolsを使う

itertools を使うことで、全ての組み合わせを生み出すイテレータを生成することができ、多重forループを避けることができます。また、生成するのはイテレータなので、全通りの組み合わせを作ったとしてもメモリを大量に消費するということもありません。作った全ての組み合わせをリストにして確認したければ、組み込み関数list()で生成したイテレータを囲むだけでOKです。（リストにした場合のメモリ消費にはご注意を）

さきほどのループを抜け出すコードをitertoolsを使って書いてみると、

itertoolsを使った場合

import itertools

all_nums_iter = itertools.product(range(10),repeat=3)

result = []
for i, j, k in all_nums_iter:
    summed = i*100 + j*10 + k
    if i == j == k and summed >= 500:
        result.append(summed)
        break
print(result)

# 実行結果 -> [555]

見やすくなったと思います。
ちなみにitertoolsには全列挙だけでなく、重複有り無しそれぞれでの組み合わせ列挙や順列の列挙もできる関数が用意されています。
※リスト内包表記でも同じことはできます。

手法B. numpy.npindexを使う

既にnumpy配列を使っている場合は、numpy.npindexを使って多重forループを避けることができます。numpy.npindexは受け取った配列の全ての値が総当たりになるようにfor文を回してくれます。下記コードではall_summed_nums.shapeが(10, 10, 10)なのでi, j, kそれぞれ0～9までの全探索for文となります。

numpy.npindexを使った場合

import numpy as np

all_summed_nums = np.array([[[100*i+10*j+k for i in range(10)] for j in range(10)] for k in range(10)], dtype="int")

result = []
for i, j, k in np.ndindex(all_summed_nums.shape):
    summed = all_summed_nums[i,j,k]
    if i == j == k and summed >= 500:
        result.append(summed)
        break
    
print(result)

# 実行結果 -> [555]

だいぶスッキリですね。

手法C. 再帰関数を使う

ちょっと難しくて使う場面も限られると思いますが、再帰関数でも全探索ができます。再帰関数がよく分からない方は再帰関数を学ぶと、どんな世界が広がるかをご覧ください。かなり分かりやすいと思います。

再帰関数を使った場合

def dfs(depth, num):
    if depth == 3:
        i, j, k = list(str(num).zfill(3))
        return num if i == j == k and num >= 500 else None
    
    for i in range(10):
        ret = dfs(depth + 1, num + 10**(2-depth) * i)
        if ret:
            return ret
        
print(dfs(0, 0))

# 実行結果 -> 555

使う場面は限られますが、特定の場面（DFS：深さ優先探索など）では直感的に分かりやすいコードを書くことができます。

参考

なぜ我々は頑なにforを避けるのか
atcoderでよく使う手法python版
再帰関数を学ぶと、どんな世界が広がるか

終わりに

最後まで読んで頂きありがとうございました。読みづらい多重forループは簡略化していきましょう〜。もし、他にも簡単に書ける方法をご存じでしたらコメントで教えて頂けると助かります。