AtCoderでPythonの勉強 その9

AtCoder Beggineres Selection 問9、ABC085C – Otoshidama です。

問題文
日本でよく使われる紙幣は、10000 円札、5000 円札、1000 円札です。以下、「お札」とはこれらのみを指します。

青橋くんが言うには、彼が祖父から受け取ったお年玉袋にはお札が N 枚入っていて、合計で Y 円だったそうですが、嘘かもしれません。このような状況がありうるか判定し、ありうる場合はお年玉袋の中身の候補を一つ見つけてください。なお、彼の祖父は十分裕福であり、お年玉袋は十分大きかったものとします。

制約
・1 ≤ N ≤ 2000
・1000 ≤ Y ≤ 2×10⁷
・N は整数である。
・Y は 1000 の倍数である。

入力は以下の形式で標準入力から与えられる。

N Y

出力
N 枚のお札の合計金額が Y 円となることがありえない場合は、-1 -1 -1 と出力せよ。
N 枚のお札の合計金額が Y 円となることがありうる場合は、そのような N 枚のお札の組み合わせの一例を「10000 円札 x 枚、5000 円札 y 枚、1000 円札 z 枚」として、x、y、z を空白で区切って出力せよ。複数の可能性が考えられるときは、そのうちどれを出力してもよい。

こういう、どうすればいいのかわからない問題は
とりあえず愚直に全探索で解いてみます。実行時間の制限は2秒、N ≤ 2000 の三重ループ。間に合うかな？

N, Y = map(int, input().split())

xans = -1
yans = -1
zans = -1

for x in range(N):
    for y in range(N - x):
        z = N - x - y
        sum = x * 10000 + y * 5000 + z * 1000
        if sum == Y:
            xans = x
            yans = y
            zans = z
print(xans, yans, zans)

二重ループでした。でも時間かかりそう。
時間短縮のため、マッチしたらループを抜けさせました。
スピードアップ改良版

N, Y = map(int, input().split())

xans = -1
yans = -1
zans = -1
match = False

for x in range(N):
    for y in range(N - x):
        z = N - x - y
        sum = x * 10000 + y * 5000 + z * 1000
        if sum == Y:
            xans = x
            yans = y
            zans = z
            match = True
        if match:
            break
    if match:
        break

print(xans, yans, zans)

入力例を試してみたら例4 (N = 2000, Y = 20000000, 正解は 2000 0 0) が -1 -1 -1 になってしまいます

あぁぁ！range(N) だと N を含まないのでした。
修正版

N, Y = map(int, input().split())

xans = -1
yans = -1
zans = -1
match = False

for x in range(N + 1):
    for y in range(N - x + 1):
        z = N - x - y
        sum = x * 10000 + y * 5000 + z * 1000
        if sum == Y:
            xans = x
            yans = y
            zans = z
            match = True
        if match:
            break
    if match:
        break

print(xans, yans, zans)

これで提出。 AC になりました。