javaで競技プログラミングをするときによく使うコードをまとめた。
あくまでも「memo」であるため、アルゴリズム自体の説明は割愛している。
C++に比べ、javaで競技プログラミングをしている人は少なく、情報が少ないため、この記事の執筆に至った。
標準ライブラリのみ使用することを想定している。
標準ライブラリについては、Javaで競技プログラミングをするときによく使う標準ライブラリで非常にわかりやすくまとめられている。
また、標準入力・標準出力が遅い等のjava固有の問題の解決には、同著者のJavaで競技プログラミングをするときの罠とテクが非常に参考になる。

基本コード

eclipseを用いる場合、 Window -> Preferences -> Java -> Code Style -> Code Templatesにて、クラスを作成した際のテンプレートを設定できる。
New Java filesにインポートするライブラリを、Method Bodyにメインメソッド内のコードを保存する。
自分は、以下のようなコードテンプレートを用意している。

import java.util.*;
public class Main {
	public static void main(String[] args) {
		Scanner sc = new Scanner(System.in);
		int N = sc.nextInt();
		System.out.println();
	}
}

また、以下のように定義している。

static int MOD = 1000000007;
static int INF = Integer.MAX_VALUE/2;

入力

配列の入力

Arrays.setAllを用いることで、for文を使うよりもすっきりとした記述ができる。

int[] a = new int[N];
Arrays.setAll(a, i -> sc.nextInt());

配列の操作

配列の最大値・最小値

for文を回さなくても、Arraysで取得可能。

Arrays.stream([配列名]).min().getAsInt()
Arrays.stream([配列名]).max().getAsInt()

配列の中央値

偶奇で場合分け。

public static int median(int[] a) {
	if(a.length % 2 == 1)
		return a[a.length/2];
	else
		return (a[a.length/2-1] + a[a.length/2]) / 2;
}

二分探索

javaではArrays.binarySearchによって二分探索が可能。
Arrays.binarySearchは配列がキーを含むならその位置を、含まないならばそのキーを挿入するべき位置を返す。
javaにはC++のlower_bound、upper_boundのようなメソッドが存在しないせず、二分探索によって個数をカウントすることができないため、自ら実装した。
このコードは二分探索を用いる様々な問題に応用可能である。

public static int lowerbound(int[] a, int key) {
	if(a[a.length-1] < key)
		return a.length;
	int lb = 0, ub = a.length-1, mid;
	do {
		mid = (ub+lb)/2;
		if(a[mid] < key)
			lb = mid + 1;
		else
			ub = mid;
	}while(lb < ub);
	return ub;
}

public static int upperbound(int[] a, int key) {
	if(a[a.length-1] <= key)
		return a.length;
	int lb = 0, ub = a.length-1, mid;
	do {
		mid = (ub+lb)/2;
		if(a[mid] <= key)
			lb = mid + 1;
		else
			ub = mid;
	}while(lb < ub);
	return ub;
}

public static int count(int[] a, int key) {
	return upperbound(a, key) - lowerbound(a, key);
}

数理的手法

桁数の取得

数字を文字列に変換すれば、その長さが桁数となる。

public static int digits(long n) {
	return String.valueOf(n).length();
}

最小公倍数

ユークリッドの互除法を再帰関数で表現する。

public static int gcd(int a, int b) {
		return b == 0 ? a: gcd(b, a % b);
}

最大公約数

最小公倍数を用いる。

public static int lcm(int a, int b) {
		return a / gcd(a, b) * b;
}

階乗 n!

再帰関数の基本である。

public static long factorial(long n){
	  return n <= 0 ? 1 : n * factorial(n-1);
}

累乗の剰余

競技プログラミングでは、オーバーフローを避けるために累乗の剰余を用いることが多い。
javaでは、これがBigIntegerクラスに実装されている。

static long modpow(int N, int K) {
	return BigInteger.valueOf(N).modPow(BigInteger.valueOf(K), BigInteger.valueOf(MOD)).longValue();
}

組み合わせ nCr

組み合わせを求める際、10^9+7で割った余りを出力する場合を考える。
Nが小さい場合(< 2000程度)は、パスカルの三角形を考え、動的計画法によって求めることができる。

int MAX = 2000;
long[][] com = new long[MAX][MAX];
for(int i = 0; i < MAX; i++)
	com[i][0] = 1;
for(int i = 1; i < MAX; i++) {
	for(int j = 1; j <= i; j++) {
		com[i][j] = com[i-1][j-1] + com[i-1][j];
		com[i][j] %= MOD;
	}
}

Nが大きい場合は、逆元を求める必要がある。
メインメソッドでCOMinit()を呼び出すことにより計算を行い、combination(n, r)によりnCrを求められるようにしている。

public static int MOD = 1_000_000_007;
public static int MAX = 100000;
public static long[] fac = new long[MAX];
public static long[] finv = new long[MAX];
public static long[] inv = new long[MAX];
	
public static void COMinit() {
	fac[0] = fac[1] = 1;
	finv[0] = finv[1] = 1;
	inv[1] = 1;
	for(int i = 2; i < MAX; i++) {
		fac[i] = fac[i-1] * i % MOD;
		inv[i] = MOD - inv[MOD%i] * (MOD/i) % MOD;
		finv[i] = finv[i-1] * inv[i] % MOD;
	}
}

public static long combination(int n, int k){
	if(n < k || n < 0 || k < 0) return 0;
	return fac[n] * (finv[k] * finv[n-k] % MOD) % MOD;
}
public static void main(String[] args) {
	COMinit();
	System.out.println(combination(2000, 3));
}

因数分解

因数分解し、HashMapに保存するメソッド。
公約数の数を求めたり、公約数を全列挙したりすることにも使える。

static Map<Integer, Integer> fact = new HashMap<>();
public static void factorization(int N){
	for(int i = 2; i <= Math.sqrt(N); i ++) {
		if(N % i == 0) {
			int n = 0;
			while(N % i == 0) {
				N /= i;
				n++;
			}
			fact.put(i, n);
		}
	}
	if(N > 1)
		fact.put(N, 1);
}

素数判定

O($\sqrt{N}$)で素数判定を行う。

static boolean isprime(int N) {
	if(N <= 1)
		return false;
	else if(N == 2)
		return true;
	for(int i = 2; i <= Math.sqrt(N); i++)
		if(N % i == 0)
			return false;
	return true;
}

素数の全列挙

エラトステネスの篩と呼ばれるアルゴリズム。
上記の素数判定を用い、O($N\sqrt{N}$)で素数の列挙を行う。
前処理として先に全ての素数判定を行っておくことで計算量を削減できる場合がある。

public static void main(String[] args) {
	int N = 100_000;
	boolean[] prime = new boolean[N+1];
	Arrays.fill(prime, true);
	prime[0] = prime[1] = false;
	for(int x = 4; x <= N; x+=2)
		prime[x] = false;
	for(int i = 3; i <= Math.sqrt(N); i+=2)
		if(prime[i])
			for(int x = i * i; x <= N; x+=i)
				prime[x] = false;
}

全順列の出力

N!通り全ての順列を出力する。

static List<String> z;
	
public static void permutation(String q, String ans){
	if(q.length() <= 1) {
		z.add(ans + q);
	 }
	 else {
	    	for (int i = 0; i < q.length(); i++) {
	    		permutation(q.substring(0, i) + q.substring(i + 1),
	    		ans + q.charAt(i));
	        }
	 }
}
public static void main(String[] args) {
	z = new ArrayList<>();
	permutation("12345","");
	for(int i = 0; i < z.size(); i++)
		System.out.println(z.get(i));
}

グラフ

木の入力

どのノードからどのノードに枝があるか、Listに保存している。

List<Integer>[] edge = new ArrayList[N];
for(int i = 0; i < N; i++)
	edge[i] = new ArrayList<>();
for (int i = 0; i < N-1; i++) {
    int a = sc.nextInt()-1;
    int b = sc.nextInt()-1;
    edge[a].add(b);
    edge[b].add(a);
}

枝に順序がある木の入力

出力が必要な場合など、枝の順番が情報として必要な場合は、枝を自作クラスとし、HashMapに保存する。
想定問題：ABC146D Coloring Edges on Tree

static int N;
static Map<Integer, HashSet<edge>> m = new HashMap<>();
static class edge{
	int to, id;
	edge(int to, int id){
		this.to = to;
		this.id = id;
	}
}
public static void main(String[] args) {
	for(int i = 0; i < N-1; i++) {
		int a = sc.nextInt();
		int b = sc.nextInt();
		if(!m.containsKey(a))
			m.put(a, new HashSet<>());
		m.get(a).add(new edge(b, i));
	}
}

グラフにおける1点からの距離

幅優先探索を用い、あるノードvからの距離を配列dに格納

int[] d = new int[N+1];
Arrays.fill(d, INF);
d[v] = 0;
Queue<Integer> q = new ArrayDeque<>();
for(int i : edge.get(v)) {
	q.offer(i);
	d[i] = 1;
}
while(q.size() > 0) {
	int x = q.poll();
	for(int i : edge.get(x)) {
		if(d[i] > d[x] + 1) {
			d[i] = d[x] + 1;
			q.offer(i);
		}
	}
}

その他手法

bit全探索

Nが小さい(N < 20程度)ときに、2^N通りを全探索する。

boolean[] x = new boolean[N];
for(int i = 0; i < 1<<N; i++) {
	for(int j = 0; j < N; j++) {
		if ((1 & i >> j) == 1)
			x[j] = true;
		else
			x[j] = false;
	}
}

union-find木

ノードのグループ分けを高速で管理することができる、union-find木の実装。
各ノードのrankを管理することによって木が深くなりすぎることを防いでいる（競技プログラミングにおいてこの高速化は多くの場合不要）。
各関数の説明は、以下。

unite:各ノードを結合
same:二つのノードが同じグループに所属しているか判定
get_size:ノードが所属するグループの大きさを取得
group_num:グループの個数を取得

public static class unionfind{
	int[] parent;
	int[] rank;
	int[] size;
	public unionfind(int n){
		this.parent = new int[n+1];
		this.rank = new int[n+1];
		this.size = new int[n+1];
		for(int i = 1; i <= n; i++) set(i);
	}
	public void set(int i) {
		parent[i] = i;
		rank[i] = 0;
		size[i] = 1;
	}
	public int find(int i) {
		if (i == parent[i])
			return parent[i];
		else
			return parent[i] = find(parent[i]);
	}
	public boolean same(int x, int y){
		return find(x) == find(y);
	}
	public void unite(int x, int y) {
		int xroot = find(x);
		int yroot = find(y);
		if(xroot == yroot)
			return;
		if(rank[xroot] > rank[yroot]) {
			parent[yroot] = xroot;
			size[xroot] += size[yroot];
		}
		else if(rank[xroot] < rank[yroot]) {
			parent[xroot] = yroot;
			size[yroot] += size[xroot];
		}
		else {
			parent[yroot] = xroot;
			size[xroot] += size[yroot];
			rank[xroot]++;
		}
	}
	public int get_size(int i) {
		return size[find(i)];
	}
	public int group_num() {
		Set<Integer> parents = new HashSet<>();
		for(int i = 1; i < this.parent.length; i++)
			parents.add(find(i));
		return parents.size();
	}
}

セグメント木

O(logN)で値の更新、および区間[a, b)の最小値の取得が可能。
init(N)で初期化、update(k, value)でk番目をvalueに更新、query(a, b)で区間[a, b)の最小値を取得。

static int N, n, INF = Integer.MAX_VALUE/2;
static int[] dat;
static void init(int N) {
	n = 1;
	while(n < N)
		n *= 2;
	dat = new int[2*n-1];
	for(int i = 0; i < 2*n-1; i++)
		dat[i] = INF;
}
static void update(int k, int a) {
	k += N-1;
	dat[k] = a;
	while(k > 0) {
		k = (k-1) / 2;
		dat[k] = Math.min(dat[k*2+1], dat[k*2+2]);
	}
}
static int query(int a, int b, int k, int l, int r) {
	if(r <= a || b <= l)
		return INF;
	else if(a <= l && r <= b)
		return dat[k];
	else
		return Math.min(query(a, b, k*2+1, l, (l+r)/2), query(a, b, k*2+2, (l+r)/2, r));
}
static int query(int a, int b) {
	return query(a, b, 0, 0, n);
}

public static void main(String[] args) {
	N = 4;
	init(N);
	int[] x = {2, 4, 3, 0};
	for(int i = 0; i < N; i++)
		update(i, x[i]);
	System.out.println(query(2, 3));
}

追記

この記事は編集途中である。
随時更新していこうと思う。