ABC051 B問題（Sum of Three Integers）を解く

AtCoder が提供しているABC（AtCoder Beginner Contest）051 B問題をC++で解いてみました。ABC051は、2017年1月7日21:00に実施されました。

AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。

B問題 Sum of Three Integers（Difficulty : 784）

問題の詳細は、リンク先をご覧ください。

ABC051 B問題 Sum of Three Integers

工夫により計算量を削減できます。AtCoder Problems による Difficulty は 784 でした。

解答案

素朴に3重ループで判定します。具体的には、$0 \leqq X, Y, Z \leqq K$ の範囲で、合計が $S$ になっているか判定します。

#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) {
			for (int z = 0; z <= k; ++z) {
				if (x + y + z == s) {
					++result;
				}
			}
		}
	}

	cout << result << endl;

	return 0;
}

残念ながらいくつかのテストケースで TLE（Time Limit Exceed）判定となります。制約 $2 \leqq K \leqq 2500$ のため、計算量は、$O(K^3) \doteqdot 1.5 \times 10^{10}$ となります。1秒では、$10^8$ 回程度の処理しかできないため、間に合いません。

C++ プログラム例（ABC051B）

3重ループでは、計算量が大きいことが分かりました。一方、$X$ と $Y$ が決まれば、$Z =S-X-Y$ である必要があります。この $Z$ が $0 \leqq Z \leqq K$ であれば、条件を満たします。この工夫で、計算量を2重ループにすることができます。

以下が、C++プログラムです。差分の行の背景色を変更しました。

#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;
			if ((0 <= z) && (z <= k)) {
				++result;
			}
		}
	}

	cout << result << endl;

	return 0;
}

AC（Accepted＝正しいプログラム）と判定されました。

最後に

簡単な工夫で3重ループを2重ループにすることができました。

引き続き ABC の問題を紹介していきます。