Java向け競プロライブラリ(素数)

素数が絡んでくる問題向けのライブラリです。

問題の制約に合わせて適宜sieveSizeの値を変更して使用します。

例えば以下の問題で役に立ちます。

AtCoder Beginner Contest 250 D - 250-like Number

AtCoder Beginner Contest 254 D - Together Square

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

final class PrimeNumber {

	private static int sieveSize = 2 * 100000;
	private static boolean[] isPrime = new boolean[sieveSize + 10];
	private static long[] cnt = new long[sieveSize + 10];

	static {
		Arrays.fill(isPrime, 2, isPrime.length, true);
		for (int i = 2; i * i <= sieveSize; ++i) {
			if (!isPrime[i]) {
				continue;
			}

			for (int j = i * i; j <= sieveSize; j += i) {
				isPrime[j] = false;
			}
		}

		for (int i = 2; i <= sieveSize; ++i) {
			cnt[i] = (isPrime[i]) ? cnt[i - 1] + 1 : cnt[i - 1];
		}
	}

	/**
	 * 与えた整数が素数ならばtrue、素数でなければfalseを返す
	 * @param n 判定対象の整数
	 * @return nが素数ならtrue、合成数ならfalse
	 */
	public static boolean isPrimeNumber(int n) {
		return isPrime[n];
	}

	/**
	 * デフォルトで指定された範囲の素数の個数を取得する
	 * @return 素数の個数
	 */
	public static long countPrimeNumbers() {
		return countPrimeNumbersLessThanOrEqualTo(sieveSize);
	}

	/**
	 * 与えた整数以下の素数の個数を取得する
	 * @param n 上限の整数
	 * @return n以下の素数の個数
	 */
	public static long countPrimeNumbersLessThanOrEqualTo(int n) {
		return cnt[n];
	}

	/**
	 * デフォルトで指定された範囲の素数を要素とする配列を取得する
	 * @return 素数を要素とする配列
	 */
	public static int[] getPrimeNumbers() {
		return getPrimeNumbersLessThanOrEqualTo(sieveSize);
	}

	/**
	 * 与えた整数以下の素数を要素とする配列を取得する
	 * @param n 上限の整数
	 * @return n以下の素数を要素とする配列
	 */
	public static int[] getPrimeNumbersLessThanOrEqualTo(int n) {
		List<Integer> p = new ArrayList<Integer>();
		for (int i = 2; i <= n; ++i) {
			if (isPrime[i]) {
				p.add(i);
			}
		}
		return p.stream().mapToInt(i -> i).toArray();
	}

	/**
	 * 与えた整数の素因数分解を実行し結果を返す
	 * @param n 素因数分解の対象となる整数
	 * @return 素因数分解の結果を格納したマップ
	 */
	public static Map<Integer, Integer> factorize(long n) {
		Map<Integer, Integer> res = new TreeMap<Integer, Integer>();
		for (int d = 2; 1l * d * d <= n; d += (d == 2) ? 1 : 2) {
			if (!isPrime[d]) {
				continue;
			}

			int cnt = 0;
			while (n % d == 0) {
				++cnt;
				n /= d;
			}

			if (cnt > 0) {
				res.put(d, cnt);
			}
		}

		if (n > 1) {
			res.put((int) n, 1);
		}
		return res;
	}

	/**
	 * 既知の素因数分解の情報をもとに約数を洗い出す
	 * @param pfs 素因数を要素とする配列
	 * @param maxPows 各素因数の個数を要素とする配列
	 * @return 約数一覧の配列
	 */
	public static long[] getDivisorsUsingFactorization(int[] pfs, int[] maxPows) {
		int[] pows = new int[pfs.length];
		List<Long> l = new ArrayList<Long>();
		findDivisorsByDfs(0, pfs, maxPows, pows, l);
		return l.stream().mapToLong(i -> i).toArray();
	}

	private static void findDivisorsByDfs(int depth, int[] pfs, int[] maxPows, int[] pows, List<Long> divDest) {
		if (depth == maxPows.length) {
			long d = 1;
			for (int i = 0, f = pfs.length; i < f; ++i) {
				for (int j = 0, fj = pows[i]; j < fj; ++j) {
					d *= 1l * pfs[i];
				}
			}
			divDest.add(d);
			return;
		}

		for (int i = 0, f = maxPows[depth]; i <= f; ++i) {
			pows[depth] = i;
			findDivisorsByDfs(depth + 1, pfs, maxPows, pows, divDest);
		}
	}

}