けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 051 B - Sum of Three Integers (200 点)

AtCoder 過去問精選10問記事だと「計算量意識するのは 300 点から」って書いてしまっているのだけど、この問題の存在がちょっと心苦しいかも ^^;

問題へのリンク

問題概要

2 つの整数  K, S が与えられます。
3 つの  0 以上  K 以下の整数  (x, y, z) の組であって、 x + y + z = S を満たすものが何通りあるか求めよ。

制約

  •  1 \le K \le 2500
  •  1 \le S \le 3K

全探索

一目見てすごく数学色強そうで怖そうなのだけど、とりあえず答えを出すコードを求めるだけなら、すべての場合を調べてしまえばよい!!!

具体的には

  •  0 以上  K 以下の整数  x
  •  0 以上  K 以下の整数  y
  •  0 以上  K 以下の整数  z

の組合せを全パターンを試して、そのうち  x + y + z = S を満たすものをカウントします。

    int count = 0;
    for (int x = 0; x <= K; ++x) {
        for (int y = 0; y <= K; ++y) {
            for (int z = 0; z <= K; ++z) {
                if (x + y + z == S) {
                    ++count;
                }
            }
        }
    }

しかし...

B 問題, 200 点問題だったら、これで通って欲しい。。。でもこれだと TLE (実行時間超過, Time Limit Exceeded) してしまう。

理由は以下の記事の 8 問目 - Otoshidama の解説に書いたのと同じなのだけど、for 文を三重にしているので、計算量は  O(K^{3}) になっている。ここで最大ケースの  K = 2500 を代入してみると  2500^{3} が大体  10^{10} くらいになる。しかし実行時間制限は  10^{9} くらいが限界なのだ。

qiita.com

改良

実は for 文を三重から二重に減らすことができます。例えば

  •  S = 10
  •  x = 2,  y = 3 の場合を考えると
  •  z 10 - 2 - 3 = 5 になるしかない

ということが見て取れます。つまり  x, y の値が決まった時点で、 z のとるべき値は決まってしまうのです。そしてそれが実際に  0 以上  K 以下になっているかどうかを確かめてあげて、それを満たしていたらカウントするようにすればよいです。

これによって、z のループが消えて、for 文が二重になりました。計算量は  O(K^{2}) に下がります!!!!!
今度こそ通ります!!!

#include <iostream>
using namespace std;

int main() {
    int K, S; cin >> K >> S;
    int count = 0;
    for (int x = 0; x <= K; ++x) {
        for (int y = 0; y <= K; ++y) {
            int z = S - x - y;
            if (z >= 0 && z <= K) {
                ++count;
            }
        }
    }
    cout << count << endl;
}