【アルゴリズム】CodeIQの「ホリエモンからの挑戦状」を解いて1,200円貰いました

気付けば1ヶ月以上もブログを放置していました。
というわけで、ちょっと前にCodeIQでやっていた、「ホリエモンからの挑戦状」というキャンペーン問題を解いて1,200円くらいもらった話をリハビリがてら書きます。

https://codeiq.jp/lp/horiemon/

と言っても企画自体は7月頭くらいだったので、もはやほとんど忘れてるんですが、2ヶ月前に自分が書いたコードを思い出しながら頑張って書きたいと思います。

問題

問題は短い棒を3本繋いで決まった長さの棒を作る場合の組み合わせパターン数を求めろというものだったのですが、単純化すると以下のようになります。

自然数Lと、N個(1≦N≦5000)の自然数とが入力値として与えられる。
N個の自然数のうち3つを選び、その和がLになるような組み合わせはいくつあるかを求めるプログラムを書け。
なお、N個の自然数の大きさは全て異なるとする。

回答晒し

で、例のごとく回答を晒します。ちょっと長いですね。

public class Main {
	
	List<Long> list = new ArrayList<Long>();
	int startIdx;
	int endIdx;
	
	public static void main(String args[]) throws IOException {
		Main self = new Main();
		self.execute();
	}
	
	private void execute() throws IOException {
		InputStreamReader isr = new InputStreamReader(System.in);
		BufferedReader br = new BufferedReader(isr);
        
		final long L = Long.parseLong(br.readLine());
		final int N = Integer.parseInt(br.readLine());
       	
		//とりあえずリスト作成
		for (int i=0; i < N; i++) {
			list.add((long) Long.parseLong(br.readLine()));
 		}
		//数の大きい順にソートしておく
		Collections.sort(list);
		Collections.reverse(list);
        
		//ここがメイン
		try {
			long ans = 0;
			for (int i=0; i<N-2; i++) {
				//i番目の要素を1つ目の要素として選択した場合の合計
				if (isLengthTooMuch(i+1, L-list.get(i), 2)) {
					break;
				}
				ans += solve2(i+1, L-list.get(i), 2);
			}
			System.out.println(ans);
		} catch (Error e) {
			System.out.println("cause=" + e.getCause());
			System.out.println("class=" + e.getClass());
			System.out.println("message=" + e.getLocalizedMessage());
		}
	}
	
	private long solve2(int idx, long length, int num) {
		long res = 0;
		
		if (!isLengthEnough(length, num)) {
			return 0;
		}
		
		//探索用endIdxの初期化
		endIdx = list.size()-1;
		startIdx = idx;
		
		for (int j=idx; j<list.size()-1; j++) {
			if (startIdx < j+1) {
				startIdx = j+1;
			}
			if (containsAfter(j+1, length-list.get(j))) {
				res++;
			}
		}
		return res;
	}
	
	private boolean containsAfter(int idx, long length) {
		if (length <= 0) {
			return false;
		}
		if (list.get(idx) < length) {
			return false;
		}
		if (list.get(list.size()-1) > length) {
			return false;
		}
		
		while (true) {
			if (Math.abs(startIdx - idx) <= 2) {
				startIdx = idx;
				break;
			} else if (list.get(startIdx) < length) {
				startIdx = startIdx-10;
			} else {
				break;
			}
		}
		
		while (true) {
			if (Math.abs(endIdx - startIdx) <= 2) {
				break;
			}
			
			int halfIdx = endIdx - Math.round((endIdx - startIdx)/2);
			long halfVal = list.get(halfIdx);
			if (halfVal == length) {
				return true;
			} else if (halfVal > length) {
				startIdx = halfIdx;
			} else {
				endIdx = halfIdx;
			}
		}
		
		for (int i = startIdx; i <= endIdx; i++) {
			if (list.get(i) < length) {
				break;
			}
			if (list.get(i).longValue() == length) {
				return true;
			}
		}
		return false;
	}

	private boolean isLengthEnough(long length, int num) {
		if (num == 1) {
			return (length >= list.get(list.size() - 1));
		} else if (num == 2) {
			return (length >= list.get(list.size() - 1) 
					+ list.get(list.size() - 2));
		} else {
			return (length >= list.get(list.size() - 1) 
					+ list.get(list.size() - 2)
					+ list.get(list.size() - 3));
		}
	}
	
	private boolean isLengthTooMuch(int idx, long length, int num) {
		if (num == 1) {
			return (length > list.get(idx));
		} else if (num == 2) {
			return (length > list.get(idx) + list.get(idx+1));
		} else {
			return (length > list.get(idx) 
					+ list.get(idx+1)
					+ list.get(idx+2));
		}
	}
}

ポイントとしては、まず最初にN個の数字を大きい順にソートしておく。
これによって、ループ処理の過程で選ばれていく組み合わせは必ず大きいものから順に並ぶことになるし、i+1番目以降の数値を使って作成できる最大の和は、最初に取れるr個の数値の和になる。
これを利用して無駄な処理を省いているのがisLengthTooMuchを呼んでいる部分で、最初の数値としてi番目のものを選んだとして、i番目とi+1番目とi+2番目の和がLに届かないなら、それ以降のループは回す必要がないよね、ということ。
これで1重目のループを回す回数を抑えます。

で、最初に選ぶのがi番目の数値ということが決まったら、次はsolve2というメソッドで、残りの要素を探す。
ここではまずisLengthEnoughというメソッドを呼んで、i番目を選んだ時点で、残りの和を最小となるような組み合わせで選んだとしてもLを超えてしまうようになっていないかどうかを事前に調査する。
リストは大きい順に並んでいるので、リストの末尾の要素2個の和がi+1番目以降の要素2つで作れる和の最小となる。
どう頑張ってもLを超えてしまうような場合はこの時点で次のループへ。
そうでない場合は、2つ目の要素としてj番目の数値を選んだとして、大きさが「L - (i番目の数値) - (j番目の数値)」の数値がj+1番目以降に存在するかどうかを調べる。
存在すれば組み合わせパターンを+1する。

で、その「大きさが「L - (i番目の数値) - (j番目の数値)」の数値がj+1番目以降に存在するかどうかを調べる」ためのメソッドがcontainsAfterなわけですが、ここでもソートしてあることが効いている。
まずi番目とj番目を選んだ時点でLを超えているようなケースはスキップするとして、次にsolve2と同様に、j+1番目以降の最大の数値（要するにj+1番目の数値）が残りの必要な長さに達していなかったスキップだし、最小の数値（要するに最後の要素）が残りの長さよりも長くてもスキップする。

で、containsAfterはさらにちょっとクセのあることをやっていて、それがwhile(true)でループしている2カ所の部分。
何をやっているかというと、擬似的な2分探索をやることで、残りの長さlengthが存在するとしたら、startIdx〜endIdxの間にしかあり得ないということを絞り込んで、実際にループを回して探す回数を抑えている。（2個目のwhile(true)）
これもリストがソートされているからできること。
しかも、このstartIdxは毎回初期化されるわけではなく、「i番目とj番目を選んだ時のstartIdxがkだったなら、i番目とj+1番目を選んだ場合のstartIdxもkと同じか、それよりちょっと手前くらいだろ」という推測のもと、2分探索自体の試行回数も少なくなるよう工夫してある。（1個目のwhile(true)）

こういったアルゴリズムで、一応合格レベルの速度は満たすことができた。
総額100万円を合格者の数で山分けという話だったので、1,200円貰えたということは合格者は800人くらいということですね。
いずれにせよ、こういう趣味プログラミングで問題解くだけでおこづかい貰えるのはハッピー。
時給換算すると800円くらいなのでちょっと辛いですが、まぁ趣味だし。数独解いてお小遣い貰えるみたいな感じだし。

おまけ

メソッド名がsolve2とかになっているのは、当然もともとはsolveがあったから。
最初はsolveを「idx番目以降の要素をnum個使って、合計lengthとなる組み合わせの数を返す」というメソッドとして定義し、再起的に呼び出すことでコーディングの省力化を図っていたのですが、上記の通り最後の1つを探す時には2分探索した方が良いとか色々あってできませんでした。
isLengthTooMuchやisLengthEnoughが無駄に汎用的に書かれているのもその名残です。
なお、C++とかの高速言語ならもしかしたら再起呼び出しのもともとのソースでも合格できたかもしれません。
そしたら時給1,000円超えるね。

以上、おわり。