AtCoder|ABC051:B - Sum of Three Integers

AtCoderの過去問対策です。ABCで解けなかった問題、ためになった問題のコードを備忘録として残します。

• 問題
• 解説
• 考え方
• 解答例1
• 解答例2
  • メモ
• デモ1
• デモ2

問題

B - Sum of Three Integers

解説

https://img.atcoder.jp/abc051/editorial.pdf

考え方

各変数を 0 から K まで回す 3 重ループだと、時間計算量は O(K³) となる。ただし、K の上限は 2500 であるため、このままでは間に合わない。 2,500³ = 15625000000(約156億)

X + Y + Z = S の式を Z = の式に変形する。

Z = S − X − Y

変数 X, Y の組を全探索して、変数 Z の値が 0 以上かつ K 以下であるような組を数える。この解法の時間計算量は、O(K²) となるため間に合う。 2,500² = 6250000(625万)

解答例1

↓TLE:実行時間制限超過でNG

#include <bits/stdc++.h>
using namespace std;

int main(){
  int K, S;
  cin >> K >> S;
  int X, Y, Z;
  int result = 0;
  
  for(X = 0; X <= K; X++){
    for(Y = 0; Y <= K; Y++){
      for(Z = 0; Z <= K; Z++){
        if(X + Y + Z == S){
          result++;
        }
      }
    }
  }
  cout << result << endl;
}

解答例2

#include <bits/stdc++.h>
using namespace std;

int main(){
  int K, S;
  cin >> K >> S;
  int result = 0;
  
  for(int X = 0; X <= K; X++){
    for(int Y = 0; Y <= K; Y++){
      int Z = S - X - Y;
      cout << X << " + " << Y << " + " << Z << " = " << S << ", K:" << K << "\n";
      if(0 <= Z and Z <= K){
        result++;
      }
    }
  }
  cout << result << endl;
}

メモ

• ループのネストを減らす

デモ1

入力

2 2

出力

0 + 0 + 2 = 2, K:2
0 + 1 + 1 = 2, K:2
0 + 2 + 0 = 2, K:2
1 + 0 + 1 = 2, K:2
1 + 1 + 0 = 2, K:2
1 + 2 + -1 = 2, K:2
2 + 0 + 0 = 2, K:2
2 + 1 + -1 = 2, K:2
2 + 2 + -2 = 2, K:2
6

デモ2

入力

5 15

出力

0 + 0 + 15 = 15, K:5
0 + 1 + 14 = 15, K:5
0 + 2 + 13 = 15, K:5
0 + 3 + 12 = 15, K:5
0 + 4 + 11 = 15, K:5
0 + 5 + 10 = 15, K:5
1 + 0 + 14 = 15, K:5
1 + 1 + 13 = 15, K:5
1 + 2 + 12 = 15, K:5
1 + 3 + 11 = 15, K:5
1 + 4 + 10 = 15, K:5
1 + 5 + 9 = 15, K:5
2 + 0 + 13 = 15, K:5
2 + 1 + 12 = 15, K:5
2 + 2 + 11 = 15, K:5
2 + 3 + 10 = 15, K:5
2 + 4 + 9 = 15, K:5
2 + 5 + 8 = 15, K:5
3 + 0 + 12 = 15, K:5
3 + 1 + 11 = 15, K:5
3 + 2 + 10 = 15, K:5
3 + 3 + 9 = 15, K:5
3 + 4 + 8 = 15, K:5
3 + 5 + 7 = 15, K:5
4 + 0 + 11 = 15, K:5
4 + 1 + 10 = 15, K:5
4 + 2 + 9 = 15, K:5
4 + 3 + 8 = 15, K:5
4 + 4 + 7 = 15, K:5
4 + 5 + 6 = 15, K:5
5 + 0 + 10 = 15, K:5
5 + 1 + 9 = 15, K:5
5 + 2 + 8 = 15, K:5
5 + 3 + 7 = 15, K:5
5 + 4 + 6 = 15, K:5
5 + 5 + 5 = 15, K:5
1