Atcoder 競技プログラミングの鉄則 演習問題集 第5問のREになってしまう理由がわからないです。

Question

実現したいこと

ソートコードの原因解説と私のコードをどのように修正すればよろしいでしょうか？

発生している問題・分からないこと

私のコードでコンパイル提出するとREになってしまいます。

該当のソースコード

N,K = map(int, input().split())
count = 0

for i in range(1, N+1):
    for j in range(1, N+1):
        for z in range(1, N+1):
            if i+j+z == K:
                count += 1

print(count)

試したこと・調べたこと

• teratailやGoogle等で検索した
• ソースコードを自分なりに変更した
• 知人に聞いた
• その他

上記の詳細・結果

特に変更したりAIは使ってません。

補足

特になし

Comments

[melian]

入力例2の場合、ループ回数は 3000**3 = 27,000,000,000 回(270億回)になります。手元のPC(Intel Core i5-8500T (6)/3.50 GHz)で実行すると10分経過しても終了しません。つまり、制限時間オーバーになっていると考えられます。なので、修正方法としては総当り(brute force)ではなく適切なアルゴリズムを採用することになります。

[TakaiY]

僕も時間がかかることが問題かと思ったのですが、その場合REでなくTLEになるんじゃないかと思って黙ってましたがどうなんでしょう。

[melian]

三重ループを二重ループにするとループ回数は900万回になります。先程の環境では1.77秒で完了なので間に合うかもしれません。もっと効率的な方法があるかと思いますが、参考にしてみてください。
https://www.mycompiler.io/view/598zz3iL7CF

[bsdfan]

時間がかかりすぎているのならREではなくTLEになると思います。
ひとつ前の質問でも、正常に動作する回答のものでREになったりしているようなので、コピペでのミスとか、その辺も注意してみたほうがいいように思います。
pythonの場合、インデントでのエラーもREになると思います。

[kikukiku]

こんな感じにすると計算量が更に半分にできそう。
https://www.mycompiler.io/view/E4Nl6x3p9Vz

[meg_]

> Atcoder 競技プログラミングの鉄則　演習問題集 第5問のREになってしまう理由がわからないです。
URLは追記できますか？

> 特に変更したりAIは使ってません。
何の話ですか？「試したこと・調べたこと」を記入する欄なのでプログラムの補足情報は本文に書く情報かと思います。

[actorbug]

REになると言っているのは、こちらの提出結果のことでしょうか。
https://atcoder.jp/contests/tessoku-book/submissions/69378422
こちらですが、Atcoder 競技プログラミングの鉄則　演習問題集 第5問ではなく、第1問に対して回答しています。
第5問に回答したいのであれば、こちらのページにソースコードを入力してください。
https://atcoder.jp/contests/tessoku-book/tasks/tessoku_book_e

[kikukiku]

>こんな感じにすると計算量が更に半分にできそう。
>https://www.mycompiler.io/view/E4Nl6x3p9Vz
上記では計算量を半分にできないことが判明しました。
すみません。取り下げます。

[TMYchan]

@actorbugさん
解答ありがとうございます。以下のリンクで再度提出しましたら、結果がTLEになりました。
しかし、上のコードだとなぜ時間がかかりすぎてしまうのですがなぜでしょうか。

[TakaiY]

> 結果がTLEになりました。
> しかし、上のコードだとなぜ時間がかかりすぎてしまうのですがなぜでしょうか。

この問題って、「1 章：アルゴリズムと計算量」の題名にもありますが、パッと思いつく方法だと計算量が多くて時間がかかりすぎてしまう問題を、適切なアルゴリズムを1つまたは複数適用することで制限時間内に解けるようになること、が目的の演習問題です。
なので、アルゴリズムや計算量の見積り方などを勉強してこの問題が解けるようになりましょう。って感じですね。

[kikukiku]

TMYchanさんのアルゴリズムは、
Nの３重ループになっているので、
最大3000×3000×3000＝270億回ループすることになります。
アルゴリズムを工夫することで
Nの２重ループにすることができます。
3000×3000＝900万回
上記にようにすれば、3000倍処理速度を上げることができるということですね。

２重ループへの仕方は、すでに他の方が回答している通りです。

[TMYchan]

TakaiYさん、kikukikuさん
返答ありがとうございます。ループ回数が多すぎて制約を超えてしまうことだと理解しました。ループの方をもう少し考えてみます。

Answer 1

質問へのコメントにも書かれていますが，提示されたコードでは実行時間の条件を満たせないと思われますので，少し視点を変えた方法を示します。

1枚目のカードが i（1 〜 N）の場合，残り（j = K - i）に着目すると以下のことが言えます。

• j が2枚目と3枚目のカードの和の範囲（2 〜 N + N）外の場合は条件を満たさないのでカウントは不要

• j が N + 1 以下の場合は，2枚目も3枚目も 1 以上が条件なので2枚目と3枚目の和が j に一致する場合の数は j - 1 通り

• j が N + 1 を超える場合は，3枚目が N 以下なので2枚目は j - N 以上の必要があり，求める場合の数は N - (j - N - 1) 通り

これらを基にした記述例を下記に示します。

また，AtCoder への 提出結果 も参照ください。

N, K = map(int, input().split())
count = 0

for i in range(1, N + 1):
    j = K - i
    if j < 2:
        continue
    elif j <= N + 1:
        count += j - 1
    elif j <= N + N:
        count += N + N - j + 1

print(count)

Comments

[kikukiku]

実行時間制限が２秒なんですね。
ちゃんと問題文を読んでいなかったです。
２秒なら１重ループにおさめておかないと無理ですね。
すばらしい。

[TMYchan]

little_streetさん
返信ありがとうございます。
上記の解説で不明点がありましたので質問失礼いたします。
iからｊに行くときに範囲が1以上から２に変わったのはなぜでしょうか？

[kikukiku]

＞j が2枚目と3枚目のカードの和の範囲（2 〜 N + N）外の場合は条件を満たさないのでカウントは不要
上記の部分のことですね。
２枚目の取りうる範囲が、１～Ｎ、
３枚目の取りうる範囲が、１～Ｎですので
jは２枚目と３枚目の和なので、２～2Nになります。

Answer 2

N,K = map(int, input().split())
count = 0

for i in range(1, N+1):
    for j in range(1, N+1):
            z = K - i -j  
            if z >= 1 and z <= N:
                count += 1

print(count)

解答に答えてくれた皆様私のほうであらたに改変をいたしたところACをとることができました。ループを三回にしないでもやる方法を編み出したので共有いたします。上記の点で悪いところがまだありましたらお伺いしてもいいでしょうか？

Comments

[kikukiku]

２重ループでもAC取れるのですね。
なら、この方法が一番アルゴリズムとしてわかりやすいと思います。
もっと処理速度は上げる方法もあるようで、
別の方が回答して頂いています。
※１重ループの方法
※０重ループの方法（階乗計算があるので０重と言えるかどうかは別問題））。

[melian]

(私から言うのも何ですが)提示されているコードは、私が9月16日の16時9分に本質問にコメントした内容と同じです。

> 三重ループを二重ループにするとループ回数は900万回になります。先程の環境では1.77秒で完了なので間に合うかもしれません。もっと効率的な方法があるかと思いますが、参考にしてみてください。
> https://www.mycompiler.io/view/598zz3iL7CF

質問にあるコードの計算量は O(N**3) で、こちらのコードでは O(N**2) になります。little_street さんの回答は O(N) になるので、処理速度を優先するのであれば O(N) の方が有利です。(実際に100倍以上高速です) 私の回答は O(1) ですが、実質的には O(N) と同等です。(comb(n, 2) を n*(n-1)//2 に置き換れば、もう少し速くなるかもしれません)

[fana]

（オーダー表記で物事を考えるなら全く意味ない話ですが）
簡便に２重ループでやるという方向だとして，ループの範囲：
> range(1, N+1)
の部分を少しだけは考えるかな？　とか思いました．

例えば「 K-i の値が 3 しかないような場合に， j のループを 1000 回とか回すのは明らかに無駄だよな…」とか．

とはいえそこで考えすぎる(?)と，既存回答の１重ループな話に行きついてしまいそうですが……
アルゴリズムが変わらない範囲（：「簡便に」の範囲）内でさくっとやれることだけはやる的な？

[kikukiku]

melianさん
＞私が9月16日の16時9分に本質問にコメントした内容と同じです。
まったくその通り。

fanaさん
＞> range(1, N+1)
＞の部分を少しだけは考えるかな？　とか思いました．
２つ目のループ中で、条件に合致したらbreakするということですよね。
それもありだと思います。
※実はコメントしていませんが、自身で検証して動作しています。
今回はACを満足するということが命題なので、そこまでしなくてもという感じですね。

[fana]

breakで抜けるのでも，ループ範囲の側を何か適切に決定するのでも良いですが，まぁそういう話です．
１つ目のループに関しても同様です（例えば K < N みたいな場合とか）．

（ぱっと見でなんか思ったことを書いちゃいましたが
計算量のオーダーが変わらないような変更はそもそも意味をなさない場面なので，無用な話でした．）

[kikukiku]

＞１つ目のループに関しても同様です（例えば K < N みたいな場合とか）
なるほど、確かにそうですね。
この条件は考えてなかった。

Answer 3

問題としては、以下の不定方程式を満たす3個の自然数(A, B, C)の組み合わせの個数を求めることになります。

A + B + C = K

制約は 1 <= A, B, C <= N です。

最初に A, B, C >= 1 の制約だけがある場合について考えます。これは Stars and Bars Algorithms を利用して解くことができます。

Stars and Bars Algorithms for Competitive Programming - GeeksforGeeks

1. Distribution among groups where groups can have 0 objects:

Imagine you have N identical objects (which we represent as stars) that you want to distribute into K distinct groups, such that each group can have 0 to N number of elements. The Stars and Bars theorem states that there are C(N+K-1, K-1) ways to do this, where C denotes the binomial coefficient.

元の不定方程式を以下の様に置き換えます。

A = a + 1
B = b + 1
C = c + 1

# a, b, c >= 0

(a + 1) + (b + 1) + (c + 1) = K
a + b + c = K - 3

この方程式(a + b + c = K - 3)を満たす a, b, c の組み合わせの個数は K - 3 個の star(星)と 3 - 1 = 2 個の bar(棒)の並べ方の個数に等しくなります。
なので、以下の二項係数(binomial coefficient)で求めることができます。

C((K - 3) + (3 - 1), 3 - 1) = C(K - 1, 2)

次に、制約 A, B, C <= N について考えます。

1. A, B, C の少なくとも1個が N + 1 以上になる場合の組み合わせの個数

先の a + b + c = K - 3 と同様に考えると、a + B + C = K - N の解の個数を求めることになります。

(a + N) + B + C = K
=>
a + B + C = K - N

C((K - N) - 1, 3 - 1) = C(K - N - 1, 2)

A, B, C のいずれかが N + 1 以上になる場合の組み合わせ個数は 3 * C(K - N - 1, 2) になるので、前述の制約 A, B, C >= 1 の組み合わせ個数から引きます。

C(K - 1, 2) - 3 * C(K - N - 1, 2)

2. A, B, C のうち2個が N + 1 以上の場合の組み合わせの個数

　1. では、A, B, C のうち2個が N + 1 以上の場合の組み合わせの個数を2回減算していますので、それを補正する必要があります。

(a + N) + (b + N) + C = K
=>
a + b + C = K - 2 * N

C((K - 2 * N) - 1, 3 - 1) = C(K - 2 * N - 1, 2)

# A, B, C から 2 個を選ぶ組み合わせ数: C(3, 2) = 3
C(K - 1, 2) - 3 * C(K - N - 1, 2) + 3 * C(K - 2 * N - 1, 2)

3. A, B, C のすべてが N + 1 以上の場合の組み合わせ個数

A, B, C のすべてが N + 1 以上の場合の組み合わせ個数についても 2. と同様に補正を行います。

(a + N) + (b + N) + (c + N) = K
=>
a + b + c = K - 3 * N

C((K - 3 * N) - 1, 3 - 1) = C(K - 3 * N - 1, 2)

C(K - 1, 2) - 3 * C(K - N - 1, 2) + 3 * C(K - 2 * N - 1, 2) - C(K - 3 * N - 1, 2)

最終的な公式は以下になります。

C(K - 1, 2) - 3 * C(K - N - 1, 2) + 3 * C(K - 2 * N - 1, 2) - C(K - 3 * N - 1, 2)

Python での実装は以下です。

from math import comb

N, K = map(int, input().split())
n = comb(K - 1, 2)
if (m := K - N - 1) >= 2: n -= 3 * comb(m, 2)
if (m := K - 2 * N - 1) >= 2:
    n += 3 * comb(m, 2)
    if (m := K - 3 * N - 1) >= 2:
        n -= comb(m, 2)
print(n)