AtCoder Beginner Contest 006 C - スフィンクスのなぞなぞ

解説を見て、なるほど！と感心したので記しておきます。

C: スフィンクスのなぞなぞ - AtCoder Beginner Contest 006 | AtCoder

自分で解いたとき

まず、この問題を見たとき、
「AtCoder Beginners Selection」の「ABC085C - Otoshidama」を思い出しました。

こちらは、価格の違うお札が3種類あり、そのお札をトータルでN枚使って、ぴったりY円にできるかという問題です。

僕が最初に解いたときは、まんまと3重for文で「TLE」を食らってしまいましたが、
実は、最後のお札の枚数は、前の2種類のお札の枚数が決まると、自動的に1つに決まるので調べる必要がなかったのです。

そして、今回の問題も、
脚の本数が異なる3種類の人間がいて、その合計人数N人で、ぴったりM本にできるか、
という問題だったので、「ABC085C - Otoshidama」と同じように解けると考えました。

提出したソースコード

import java.util.Scanner;
 
public class Main {
	public static void main(String[] args) {
		Scanner sc = new Scanner(System.in);
		int n = sc.nextInt();
		int m = sc.nextInt();
		// 大人の足のループ変数
		for (int i = 0; i <= n; i++) {
			// 老人の足のループ変数
			for (int j = 0; j <= n - i; j++) {
				int k = n - i - j; // 赤ちゃんの足
				if (2 * i + 3 * j + 4 * k == m) {
					System.out.println(i + " " + j + " " + k);
					System.exit(0);
				}
			}
		}
		System.out.println("-1 -1 -1");
 
	}
 
}

結局これで「TLE」をもらってしまうのですが、
上記の「ABC085C - Otoshidama」と違うのは、制約のNの量です！

「ABC085C - Otoshidama」で入力で受け取る枚数Nは、最大 2000 でしたが、
今回の「C - スフィンクスのなぞなぞ」では、最大で 10^5 受け取ります。

同じように二重for文で実行した場合、そのループ回数は最大で、
・「ABC085C - Otoshidama」：2,000^2 = 4,000,000回（4 * 10^6）
・「C - スフィンクスのなぞなぞ」：100,000^2 = 10,000,000,000回（10^10）
となります。

ここで、
1 秒間で処理できる for 文ループの回数は、10^8 = 100,000,000 回程度

ということから、今回の問題は「ABC085C - Otoshidama」とは
同じ方法では解けないということが分かりますね。

qiita.com
（1秒間に処理できるループの回数のお話はこちら。）

ではどう解くのか？

前述の通り、今回与えられる入力値は「10^5」と大きな値であるため、
オーダーがN^2になることは許されません。

「解説」によると、この問題の解法には、

・つるかめ算を使う
・老人の数を固定する

の2パターンあるみたいです。
どちらもオーダーがNで済むように工夫されています。

そして、僕が個人的に感動したのは、後者の「老人の数の固定」です。

つるかめ算

恥ずかしながら、つるかめ算は名前しか知らなかったので、
さっきググって覚えました。
www.manabinoba.com

import java.util.Scanner;

public class Main {
	public static void main(String[] args) {
		Scanner sc = new Scanner(System.in);
		int n = sc.nextInt();
		int m = sc.nextInt();
		// 老人（3本足）の足
		for (int i = 0; i <= n && m >= i * 3; i++) {
			int j = 0; // 2本足の数
			int k = 0; // 4本足の数
			// つるかめ算
			// 残りすべてが2本足の場合の、足の総数を数える
			int needLegs = m - i * 3;
			int numJ = 2 * (n - i);
			if (numJ >= needLegs) {// 本数が上回った場合
				int over = numJ - needLegs;
				if (over % 2 == 0) {
					k = over / 2;
					j = (numJ / 2) - k;
				} else { // 奇数なら2本足と4本足で賄えないので飛ばす
					continue;
				}
			} else if (numJ < needLegs) {// 本数が足りない場合
				int more = needLegs - numJ;
				if (more % 2 == 0) {
					k = more / 2;
					j = (numJ / 2) - k;
				} else {
					continue;
				}

			}

			if (3 * i + 2 * j + 4 * k == m && j >= 0 && k >= 0) {
				// 出力する順は、大人(j)→老人(i)→赤ちゃん(k)
				System.out.println(j + " " + i + " " + k);
				System.exit(0);
			}
		}
		System.out.println("-1 -1 -1");

	}

}

つるかめ算を使うことで、
老人の人数を決めたときに、大人と赤ちゃんの人数も決められるので、
必要な計算量は、老人が0人からN人までの計N+1回だけで済みました。

老人の数の固定

2本足の大人、3本足の老人、4本足の赤ちゃんがいたとき、
老人2人は計6本、大人と赤ちゃん一人ずつの2人組も計6本となります。

人数がどちらも2人で変わらず、足の本数も変わらないので、
2人の老人は、すべて大人と赤ちゃんのペアによって置き換えることができます。

例）老人50人は、大人25人と赤ちゃん25人で置き換えられる。
例）老人11人のうち10人は、大人5人と赤ちゃん5人で置き換えられる。

よって、調べる老人の人数というのは、
0人の場合と、1人の場合の2パターンしかないのです！

import java.util.Scanner;

public class Main {
	public static void main(String[] args) {
		Scanner sc = new Scanner(System.in);
		int n = sc.nextInt();
		int m = sc.nextInt();
		// 老人2人=6本足、大人1人赤ちゃん1人の計2人=6本足
		// 人数と足の本数が変わらないので、老人は0と1人のケースだけ調べれば良い
		for (int a = 0; a <= 1; a++) {
			for (int b = 0; b <= n - a; b++) {
				int c = n - a - b;
				if (3 * a + 2 * b + 4 * c == m) {
					System.out.println(b + " " + a + " " + c);
					System.exit(0);
				}
			}
		}
		System.out.println("-1 -1 -1");

	}

}

外側のループで、老人を0から1人までの2通り調べ、
内側のループで、大人の数を0からNまで調べます。
赤ちゃんの人数は「ABC085C - Otoshidama」と同じく、残った人数から自動的に決まります。

すると、オーダーは2*(N+1)となるので、10^8の範囲に収まります。

こういった、ひらめきや工夫で計算量を減らすタイプの解説を見ると、
「おわぁ〜〜！なるほどぉ〜〜！！！！！」ってすごく感心させられます。

この考え方を知っていれば、
もし今後、4本足、5本足、7本足の得体の知れない生き物が現れたときなんかに、

5本足が3匹は、4本足が2匹と7本足が1匹の計3匹で、同じように補えるので、
5本足は0匹、1匹、2匹の3パターンだけ調べれば良いことが分かりますね。

忘れないうちに、そんな問題お待ちしてます。