AtCoder Beginner Contest 051のB問題をコーディングの練習がてらに解いてみました。
AtCoder Beginner Contestは競技プログラミングの初級者程度を想定したコンテストで、A問題、B問題、C問題くらいまではちょうどいい難易度になっています。(D問題以降は難しい)
問題
2 つの整数 K, S が与えられます。
3 つの変数 X, Y, Z があり、0 ≦ X, Y, Z ≦ K を満たす整数の値を取ります。
X + Y + Z = S を満たす X, Y, Z への値の割り当ては何通りありますか。
制約
- 2 ≦ K ≦ 2500
- 0 ≦ S ≦ 3K
- K, S は整数である。
問題や制約、入力、出力に関しては以下のリンクから参照してください。
https://atcoder.jp/contests/abc051/tasks/abc051_b
解答例
言語はKotlinで解きました。
1 2 3 4 5 6 7 8 9 10 11 12 13 |
fun main() { val (K, S) = readLine()?.split(" ")?.map { it.toLong() } ?: throw Exception() var count = 0 for (x in 0 until K + 1) { for (y in 0 until K + 1) { val z = (S - (x + y)) if (z in 0..K) { count++ } } } println(count) } |
解説
まず問題を読んで思いつくのがfor文を用いて3重ループで解く方法です。
しかし、3重ループで解くと時間がかかってしまい、タイムアウトしてしまいます。
そこで、3重ループ目のZを求める際にループを使わずに求めることによって、タイムアウトせずに解きたいと思います。
具体的にはSからXとYを足した値を引いたZを計算し、そのZがK以下の時、0 ≦ X, Y, Z ≦ KとX + Y + Z = S を満たすので、2重ループで解けるようになります。
終わりに
AtCoderのAtCoder Beginner Contestはロジックを書く練習にちょうど良いので、これからも定期的に練習をして、C問題までは確実に解けるようになりたいと思います。
参照
https://atcoder.jp/contests/abc051/tasks/abc051_b