プログラミング
記事内に商品プロモーションを含む場合があります

素因数分解(pfs)の繰り返しを高速化する【Go言語】

Aru

競技プログラミングでは、素因数分解を何度も繰り返す問題が出題されることがあります。 基本的な素因数分解のアルゴリズムは、1つの数に対して $O(\sqrt{N})$ の計算量がかかります。数回の計算であれば問題ありませんが、これを数万回、数十万回と繰り返すと、Go言語といえども 時間超過 (TLE) してしまうことがあります。

この記事では、何度も素因数分解を繰り返す場合に効率の良い「最小素因数(SPF: Smallest Prime Factor)」を用いたテンプレートを紹介します。

素因数分解のコード(典型)

結果をMapで返す関数

まずは、一般的によく使われる「試し割り」による素因数分解の実装です。$2$
以上の整数 $N$ を受け取り、素因数をキー、その指数を値とするマップを返します。

素因数分解については、よく見かけるコードだと思いますので説明は省略します。

// PfsMap : 素因数分解し、マップを作成
func PfsMap(n int) map[int]int {
	pfs := make(map[int]int)
	for n%2 == 0 {
		pfs[2] = pfs[2] + 1
		n = n / 2
	}

	for i := 3; i*i <= n; i = i + 2 {
		for n%i == 0 {
			pfs[i] = pfs[i] + 1
			n = n / i
		}
	}

	if n > 2 {
		pfs[n] = pfs[n] + 1
	}

	return pfs
}

計算量

このアルゴリズムの計算量は、入力 $N$ に対して $O(\sqrt{N})$ です。

$2$ で割れるだけ割った後、$3$ から $\sqrt{N}$ までの奇数で順に割り切れるか試していきます。

最悪ケース($N$ が素数の場合)では $\sqrt{N}$ までループが回ります。

くり返す場合の問題

例えば、制約が $N \le 10^6$ で、クエリ数(素因数分解する回数)が $Q = 10^5$ の場合を考えてみます。

$$10^5 \times \sqrt{10^6} = 10^5 \times 10^3 = 10^8$$

全体で約 $10^8$ 回程度のループ処理が必要になります。Go言語は高速ですが、剰余演算などのコストもあり、制限時間が厳しい問題(2秒など)では TLE する可能性が高くなります。さらに $N$ が $10^9$ などの場合、この手法では確実に間に合いません。

繰り返し素因数分解する場合の関数

多数の数値に対して素因数分解を行う場合、「エラトステネスのふるい」を応用した前計算による高速化が有効です。

この手法では、あらかじめある整数までの「最小素因数(SPF)」を計算し配列に記録しておきます。この配列を使うことで、素因数分解のクエリを高速化するのがこの手法です。

コード

コードは以下の通りです。

const blockSize = 100000

var spf []int

func FastPfsMap(n int) map[int]int {
	// 現在のサイズで足りない場合のみ拡張
	if n >= len(spf) {
		oldSize := len(spf)
		// 必要なサイズを計算(nを含む次のブロックサイズ倍数)
		targetSize := ((n / blockSize) + 1) * blockSize

		// 1. 配列を拡張
		newPart := make([]int, targetSize-oldSize)
		spf = append(spf, newPart...)

		// 2. 新しい部分を初期化(自分自身を最小素因数とする)
		for i := oldSize; i < targetSize; i++ {
			spf[i] = i
		}

		// 3. ふるい落とし
		// 素数i=2から開始し、sqrt(targetSize)までチェック
		for i := 2; i*i < targetSize; i++ {
			if spf[i] == i {
				// ループの開始位置(start)を決定
				// 基本は i*i からだが、拡張時は「拡張領域の先頭」までスキップする
				start := i * i
				if start < oldSize {
					// 拡張領域(oldSize)以上になる、最小の i の倍数を計算
					// (oldSize + i - 1) / i は切り上げ除算です
					start = ((oldSize + i - 1) / i) * i
				}

				// 拡張領域内だけを更新
				for j := start; j < targetSize; j += i {
					if spf[j] == j {
						spf[j] = i
					}
				}
			}
		}
	}

	// 素因数分解処理
	pfs := make(map[int]int)
	for n > 1 {
		p := spf[n]
		pfs[p]++
		n /= p
	}

	return pfs
}

このコードでは、引数 n が前計算しておいた範囲を超える場合には、blockSize 単位で前計算しておく範囲を拡張する処理を行います。これにより、初期化関数を呼び出さなくても自動で、拡張が行われて便利です。

このコードは並列実行を考慮していないため、goroutine から呼び出す場合、初期化が重複して実行されてしまう問題があります。sync や lock を用いて排他制御する必要があります。

競技プログラミング用途では通常単一スレッドで実行されるため問題になりませんが、ライブラリ用途では注意が必要です。

テーブルを利用した素因数を求める部分は、以下の部分で、ほとんどが前計算になります。

素因数分解の部分
	// 素因数分解処理
	pfs := make(map[int]int)
	for n > 1 {
		p := spf[n]
		pfs[p]++
		n /= p
	}

	return pfs

n が割れる数が spf[n] に格納されているので、これで割る処理を繰り返すだけです。

素因数分解の結果をmapに格納していますが、少し変更してスライスにすることも可能です。mapより高速なので、さらに速度が必要な場合はおすすめです。

計算量(繰り返し行う場合)

この手法の最大の利点は、1回あたりのクエリ計算量が $O(\log N)$ になることです。

  1. 前計算(配列構築)
    エラトステネスのふるいと同様のロジックで、最大値 $M$ まで作成する場合、約 $O(M \log \log M)$ です。これは非常に高速です。上記のコードでは必要に応じて拡張するため、無駄な領域を確保しません。
  2. クエリ(素因数分解)
    $N$ を最小素因数で割り続けるだけなので、素因数の個数分しかループしません。これは $N \le 10^9$ でも最大30回程度です。

つまり、典型コードと比較すると以下のようになります。

手法1回あたりの計算量クエリ Q 回の合計特徴
典型(試し割り)$O(\sqrt{N})$$O(Q\sqrt{N})$実装が楽、メモリを使わない
高速化(SPF)$O(\log N)$$O(M \log \log M + Q \log N)$圧倒的に速い、メモリを消費する

$N \le 10^6, Q = 10^5$ の場合でも、クエリ部分は一瞬で終わります。

要素が多い場合に、素因数を求める問題が出題された場合にはこのコードが役に立つかと思います。

AtCoderのABCの場合、値の範囲が事前計算してメモリに格納できる(10^7)程度になっているので、「ここからこの手法が使える」ことを類推できたりします。

まとめ

まとめです。以下のような使い分けをすれば良いかと思います。

  • 単発の素因数分解や $N$ が非常に大きい($10^{12}$など)場合は、従来の $O(\sqrt{N})$ の手法を使います。
  • $N \le 10^6 \sim 10^7$ 程度で、何度も素因数分解が必要な場合は、今回紹介した FastPfsMap を使うと TLE を回避できます。

問題の制約($N$ の最大値とクエリ回数)を見て、適切なアプローチを選択できるようにしておきましょう。

メールアドレスが公開されることはありません。 が付いている欄は必須項目です

ABOUT ME
ある/Aru
ある/Aru
IT&機械学習エンジニア/ファイナンシャルプランナー(CFP®)
専門分野は並列処理・画像処理・機械学習・ディープラーニング。プログラミング言語はC, C++, Go, Pythonを中心として色々利用。現在は、Kaggle, 競プロなどをしながら悠々自適に活動中 保有資格:CFP, マンション管理士、管理業務主任、宅建士など
記事URLをコピーしました