B - Sum of Three Integers

コメント歓迎です。
作者の出題意図にまんまと引っ掛かった問題です。

問題

AtCoder 版！蟻本 (初級編) - Qiita より

beta.atcoder.jp

方針1

3つの変数X, Y, Zの和がSとなるパターン数をカウント。
for文でX, Y, Zの値を変えながら全パターンの和を計算。
計算するごとに和がSかどうか判定し、該当したらカウンターに+1する。

0 ≦ X, Y, Z ≦ K の制約より、計算量としてはとなります。
結論として、この方針では計算量が多すぎるため実行時間制限をクリアできません。
(Time Limit Exceeded)

コード1(TLE)

#include <iostream>
using namespace std;
 
int main(){
  // 入力
  int K, S;
  int count = 0;

  cin >> K; // X, Y, Zの最大値
  cin >> S; // 和
 
  // 全探索 計算量K^3
  for (int X = 0; X < K + 1; ++X) {
    for (int Y = 0; Y < K + 1; ++Y) {
      for (int Z = 0; Z < K + 1; ++Z) {
        // 和がSと一致したらカウンターに+1
        if (X + Y + Z == S) {
          ++count;
        }
      }
    }
  }
  // 出力
  cout << count << endl;
}

方針2

実行時間制限をクリアできなかったため計算量を減らす工夫を考えます。

まず思いついた方法は、X+Y>Sの場合はZの計算をスキップする方法です。
X+Y>Sの場合は、Zは0以上の制約があるためZを加算してもX+Y+Z=Sの条件を満たせません。

もう一つの方法として、X+Y+Z>Sの場合はそれ以上値が大きくなる計算をスキップする方法です。

この二つの方法で計算量を減らすことができますが、減少量としてはK, Sの条件に依存します。
Sが小さい値であれば効果大だと思います。

しかし、結果としてはこの方針でも実行時間制限をクリアできません。

コード2(TLE)

#include <iostream>
using namespace std;
 
int main(){
  // 入力
  int K, S;
  int count = 0;

  cin >> K; // X, Y, Zの最大値
  cin >> S; // 和
 
  // 全探索 計算量K^3
  for (int X = 0; X < K + 1; ++X) {
    for (int Y = 0; Y < K + 1; ++Y) {
      // X + Y > SならZの計算をスキップ
      if (X + Y > S) break;
      for (int Z = 0; Z < K + 1; ++Z) {
        // 和がSと一致したらカウンターに+1
        if (X + Y + Z == S) {
          ++count;
        }
        // X + Y + Z > Sならそれ以上Zを増やして計算しない
        else if (X + Y + Z > S) break;
      }
    }
  }
  // 出力
  cout << count << endl;
}

方針3

実行時間制限をクリアできなかったためさらに計算量を減らす工夫を考えます。
小細工ではクリアできないのでは？と感じたので根本的に見直そうと考えました。

改めて考えると、XとYが決まっているとき条件を満たすZはひとつしかありません。
例えば、S=10でX=1, Y=2のとき、条件を満たすZは7(S-X-Y)しかありません。
この7が0 ≦ Z ≦ Kを満たしていれば良いわけです。

そうなると、for文でループさせる変数はX, Yのみで計算量としては で済みそうです。

結果としてはこの方針でクリア(ACcept)できました。
コードはシンプルになりました。

コード3(AC)

#include <iostream>
using namespace std;
 
int main(){
  // 入力
  int K, S;
  int count = 0;

  cin >> K; // X, Y, Zの最大値
  cin >> S; // 和
 
  // 全探索 計算量K^2
  for (int X = 0; X < K + 1; ++X) {
    for (int Y = 0; Y < K + 1; ++Y) {
      // 0 <= S - X - Y <= Kならカウンターを+1
      if (0 <= S - X - Y && S - X - Y <= K) ++count;
    }
  }
  // 出力
  cout << count << endl;
}