Pythonで単純な全探索を実装してみる-ABC051

全探索

全探索にはbit全探索やDFSなどいろいろありますが、今回は最も簡単な基本の全探索をPythonで実装していきます。

問題

問題は ABC 051 B-Sum of Three Integersを扱います。

B - Sum of Three Integers

問題文

2 つの整数 K,S が与えられます。
3 つの変数 X,Y,Z があり、0≦X,Y,Z≦K を満たす整数の値を取ります。
X+Y+Z=S を満たす X,Y,Z への値の割り当ては何通りありますか。

制約

2≦K≦2500
0≦S≦3K
K,Sは整数

実装

基本の全探索はfor文を2重3重にして回すことが多いです。(4重はあまり使わないかな)

この問題では0≦X,Y,Z≦Kの範囲でX+Y+Z=Sを満たすものを数え上げればいいので範囲内を全探索します。

k,s = map(int,input().split())
ans = 0
for x in range(k+1):　#範囲がk以下なのでrangeはk+1
    for y in range(k+1):
        for z in range(k+1):
            if x+y+z == s:
                ans += 1
print(ans)        

このコードで入力例の答えは正しく出力されましたが、いざ提出してみると……

TLEがでました。

3重のfor文では間に合わないのでなんとか実行時間を短くしなければなりません。

そこで X+Y+Z=S の条件に注目します。この条件をよく見てみると、XとYが決まればSは与えられているので必然的にZの値が決まることがわかります。
先ほどのコードではXとYが決まっているのにも関わらずZをfor文を使って探索しています。

そこでZのfor文を消し、以下のように書き換えました。

k,s = map(int,input().split())
ans = 0
for x in range(k+1):
    for y in range(k+1):
        z = s-(x+y)
        if z >= 0 and z <= k　#zの範囲を間違えないように注意
            ans += 1
print(ans)        

今回はちゃんと通りました。

基本の全探索を実装するのは簡単ですが、計算時間を短くするためには少し考えないといけませんね。

この問題の他にABC 085 C-Otoshidamaも全探索の練習問題としていいと思うのでぜひ解いてみてください。

おわり