AtCoder ABC 051 B – Sum of Three Integers Python解説

2021年11月15日2021年11月21日

Sum of Three Integers

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

AtCoder株式会社　AtCoder Beginner Contest 051

x、y、zでそれぞれfor文でループさせ、x + y + z = sを満たす組み合わせをカウントしていくのが基本的な解法になりますが、k<=2500まで大きくなりうるのでfor文を3重ループさせると時間が間に合いません。

ですので、条件x + y + z = sより、zは z = s – (x + y)で置き換え、3重ループを2重ループに減らします。
これで制限時間に間に合います。

k, s = map(int, input().split())

cnt = 0

for x in range(k+1): # *1
    for y in range(k+1):
        z = s - (x + y)
        if 0 <= z <= k:
            cnt += 1
            
print(cnt)

(*1) 今回はkの数まで使いたいのでfor文はk + 1までにしています。for文でxとyの組み合わせを出し、zが0以上k以下の条件を満たしていたらカウントアップしています。