二分探索について

今回もAtCoderの勉強です。
昨日の記事で書くと言っていた二分探索について今回は書きます。

ちなみに、昨日のABC166の結果はこちらです。
developer-bpeldi2oerkd8.hatenablog.com

二分探索とは

全体を2つに分けながら条件を満たすものを見つける探索方法です。
競プロの基本は全探索ですが、二分探索では全探索よりも高速に見つけることができます。
計算量はO(logN)とかなり高速に探索できます。

ただし、事前に探す対象の配列をソートしておく必要があります。
JavaではソートにArrays.sort()を用いることが多く、この計算量はO(NlogN)となるので、ソートされていない配列に対しては全体でO(NlogN)で見つけることができます。

どのように解くのか（Javaの場合）

二分探索の方法にはいろいろありますが、ここでは競プロ界隈で推奨されている「めぐる式二分探索」を使います。
めぐる式二分探索の説明はこちらを参考にさせていただきました。
qiita.com

いくつか実装方法を紹介しているサイトはあったのですが、いまいち自分にヒットするものがなかったので、元のサイトを参考にしながら、自分で実装してみたのがこちらです。
（バグが残っているかもしれません）
github.com

簡単に自分で実装したものの説明をします。
基本的にはisOKの条件を変えることで対応できます。

ただし、binary_searchメソッドの引数を一部変えています。

まず、配列をラッパークラスに変更しました。
これは、配列を逆順にソートしたものを引数として入れられるようにするためです。

次に、ngとokの引数を追加しました。
これは、配列の中の範囲を指定して探索できるようにするためです。
ng<okのとき、全体から探索したい場合は、ng=0, ok=a.lengthとします。（右側が真の場合）
ng>okのとき、全体から探索したい場合は、ng=a.length, ok=0とします。（左側が真の場合）

実装した自作のライブラリを使って、二分探索の問題を解いてみました。
解いた問題は、ABC077のC問題 Snuke Festival です。
atcoder.jp

この問題の肝は普通に全探索をAから順にやるとTLEしてしまうので、二分探索を使って計算量を減らしてやることです。
さらに、A→B→CとみていくとO(N^2logN)となってしまうので、Bを基準にしてAとCを二分探索していきます。
（全く気づきませんでした・・・）

import java.util.Arrays;
import java.util.Scanner;

public class Main {
	//java11

	public static void main(String[] args) {

		Scanner sc = new Scanner(System.in);

		int N = sc.nextInt();

		Integer[] A = new Integer[N];
		Integer[] B = new Integer[N];
		Integer[] C = new Integer[N];

		for(int i=0; i<N; i++) {
			A[i] = sc.nextInt();
		}

		for(int i=0; i<N; i++) {
			B[i] = sc.nextInt();
		}

		for(int i=0; i<N; i++) {
			C[i] = sc.nextInt();
		}

		Arrays.sort(A);
		Arrays.sort(B);
		Arrays.sort(C);

		long ans = 0;
		for(int i=0; i<N; i++) {
			int key = B[i];
			int a_n = binary_search(A, key, 0, A.length);
			if(a_n <= 0 || a_n > A.length)
				continue;

			key++;
			int c_n = binary_search(C, key, 0, C.length);
			if(c_n < 0 || c_n >= C.length)
				continue;

			ans += (long)a_n * (N-c_n);
		}

		System.out.println(ans);

	}

	//okの条件（ここを問題ごとに変える）(int)
	public static boolean isOK(Integer[] a, int index, int key) {
		if(a[index] >= key)
			return true;
		else
			return false;
	}

	//めぐる式二分探索(デフォルトはkey以上を満たす最小のindexを返す)(int)
	public static int binary_search(Integer[] a, int key, int ng, int ok) {
//		int ng = -1;
//		int ok = a.length;
		if(ng < ok)
			ng--;
		else
			ok--;

		while(Math.abs(ok-ng) > 1) {
			int mid = (ok+ng) / 2;

			if(isOK(a, mid, key))
				ok = mid;
			else
				ng = mid;
		}

		return ok;
	}
}

ということで、二分探索に関する記事でした。
緑diffの問題が解けるようにこれからも勉強を続けていきます。