【AtCoder版！蟻本】ABC051 B - Sum of Three Integers【準備編】

AtCoder版！蟻本の準備編で類題としてあげられているABC051 BをPythonで解いていきます。

リンク

atcoder.jp

AtCoder Beginner Contest 051 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 は整数である。

解答例

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:
            ans += 1

print(ans)

解説

k以下の正の整数x,y,zがあり、x + y + z = sを満たす組み合わせの数を答える問題です。

x,y,zの3重ループにすると計算量が多くなりますので、xとyの2重ループになるよう工夫が必要です。
x + y + z = s から z = s - x - y と式変形できます。
なので、xとyが決まれば候補となるzが求まり、この候補が条件を満たすか（正の整数でk以下か）を判定します。
条件を満たせばansに1を足していけば、ループを抜けたあとのansが求める解答です。

AtCoder版！蟻本（初級編）の記事リンクページはこちら
ebisuke33.hatenablog.com

リンク