電電のブログ

電電生だったひとのぼやき

Python3 での多重ループ文の抜け方

ほい電電です.
今回はちょっと技術っぽいことを書いていきます

まあといってもシステム的なところではなく,もっと単純なところの説明です.
これに至ったのはどこかのブログでfor else文の存在を知ったからなんですが,元のブログがわからなくなったのでまとめついでに書きます.


さてみなさんPython3書いてますか?
僕は最近競技プログラムにはまっていて,気づいたら問題解いててちょっとやばいです.研究したいんですが,提携先のデータ譲渡が行われてなくて,どうしようもないという悲しい状況です.
(オープンデータで独自にやる路線が有力になりつつありとても辛い...)
で話は戻りまして,競技プログラミングしてるとちょくちょく二重や三重ループ文を書くのですが,
ループから抜け出すのがなかなかめんどくさいんです.

そこで見つけたfor else文がなかなか面白かったので備忘録として書いておきます.
ドキュメントもあるので「俺は一時ソースしか見ないぜ」って人は下記リンクから飛んでください.
docs.python.org


ここからは備忘録です.
まず注意として,多重ループをする際には計算量に注意しましょう.単純な繰り返し処理ですが計算のオーダーは絶対に大きくなるのでそもそもの計算量が少ない,or すぐ抜けれるはずなどの目論見が立ってるときに使いましょう.
競技プログラミングにおいて,計算オーダーが 10**6 を超え出すあたりからは大体問題のアプローチが間違っています.グラフにするなり,部分和を作るなり,工夫すべきところを考えましょう.(自戒)

本題に入ります.多重ループを抜けるにはいくつか方法が考えられます.
ここでは例として以下のような問題を考えます.
自然数M,N,Xが与えられる.この数字のそれぞれについて1以上M,N以下の数字をそれぞれm,nとする.このm,nについてm**nがXの倍数となるようなものが存在すればその数字を,存在しなければ"-1"を返せ』

1\leq M,N\leq 100\\ 1\leq X \leq 10^5

テストケース1
M,N,X = 5,3,121


テストケース2
M,N,X = 5,7,300


テストケース3
M,N,X = 3,4,81


このような問題を解いていきましょう.
ここでは問題設定のM,Nが小さいので全探索でも間に合いそうです.(多分もっと簡単に解けると思います.)
全探索するためには二重ループが必要なのですが,以下のような方法があります.

  1. flag を立てる
  2. 関数化する
  3. for else文を使う

1.flag を立てる

これが一番最初に浮かぶ人が多いと思います.ただループが深くなりすぎたり,flagを立てすぎたりすると混乱したりします.(今回の問題も実はフラッグをわざわざ作る必要はなくて,ans!=-1で代用できます.)

N,M,X = map(int,input().split())
flag = False
ans = -1
for n in range(1,N+1):
    for m in range(1,M+1):
        # print(n,m)
        if m**n % X == 0:
            ans = m**n
            flag = True
            break
    if flag:
        break

print(ans)

2.関数化する

これも結構多いと思います.
関数化することでreturn 文で値を返すことができ,その上残りのfor の処理をしなくていいという利点があります.ただ問題としては関数で使用するものが多いと入力値が少々めんどくさいです.これは全部グローバル関数にすればいいという力技で解決(???)できますが,競技プログラミング限定の話だと思います.
実際に実用化する製品とかサービスでグローバル関数をたくさん定義するとひどいことになるのでやめましょう.

N,M,X = map(int,input().split())
flag = False
ans = -1

def solve(N,M,X):
    for n in range(1,N+1):
        for m in range(1,M+1):
            print(n,m)
            if m**n % X == 0:
                ans = m**n
                return ans

ans = solve(N,M,X)

print(ans)

3.for else文を使う

これを思いつく人は少ないと思います.僕はそもそもfor文にelseが使えること自体知らなかったです.
でこのelseなんですが,if elseとは違って,for文が全ての反復処理を使い切ってループが終了すると実行されます.つまり,forの中でbreakが実行されてループが途中で終了した時実行されません.
これを利用すると以下のように書き直すことができます.

N,M,X = map(int,input().split())
flag = False
ans = -1

for n in range(1,N+1):
    for m in range(1,M+1):
        print(n,m)
        if m**n % X == 0:
            ans = m**n
            break # break 1
    else: 
        continue
    break # break 2   
print(ans)

ちょっと綺麗ですね.
解説をします.上部のbreakをbreak 1,下部のbreakをbreak 2とします.
この時,for else文はfor 文がbreakなしで回りきった時に実行されます.ここではcontinueが実行されるのでさらにその下のbreak2がスキップされます.一方で,上部のbreak1が実行されるとforが回りきっておらず,else以下は実行されません.そのため,下部のbreak2が実行されて,多重ループを抜ける.という動きになります.
これなら3重ループとかでも書きやすい気がします.
まあそもそも知ってる人があまり多くないのであまり使うなって意見もあるかもしれませんが...


今日はここまでです.ほなまた