ただの技術メモ

個人備忘録

高速な素数判定

競技プログラミングで度々出題される素数判定についてまとめる。

単純な素因数分解問題

以下のような単純に素因数分解する問題を考える。

https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=NTL_1_A&lang=ja

2から√nまでの整数で割っていき、割り切れれば割った整数を素数として追加し、最終的に残った数字が1より大きければ最後の素数になるのでその数も素数として追加する。

コードで表現すると以下のようになる。

package main

import "fmt"

func main() {
    var n int
    fmt.Scan(&n)

    fmt.Printf("%d:", n)

    primeFactor := make([]int, 0)
    for i := 2; i*i < n+1; i++ { // 2~√nまでの約数を考える
        for {
            if n%i == 0 {
                primeFactor = append(primeFactor, i)
                n /= i
            } else {
                break
            }
        }
    }
    if n > 2 { // 1以上であれば最後の約数(素数)
        primeFactor = append(primeFactor, n)
    }

    for i := range primeFactor {
        fmt.Printf(" %d", primeFactor[i])
    }
    fmt.Println()
}

通常の素因数分解はこのように解ける。

エラトステネスの篩を使った素数判定

以下の問題を考える。

atcoder.jp

以下のコードのように、問題文通りにlからrまでの奇数xについて素数かどうかを判定するとTLEしてしまう。

package main

import "fmt"

func main() {
    var Q int
    fmt.Scan(&Q)
    l := make([]int, Q)
    r := make([]int, Q)
    for i := 0; i < Q; i++ {
        fmt.Scan(&l[i], &r[i])
    }

    for i := 0; i < Q; i++ {
        var count int
        for x := l[i]; x < r[i]+1; x += 2 {
            if x == 1 { // 1は素数ではないので除外
                continue
            }
            if isPrime(x) && isPrime((x+1)/2) {
                count++
            }
        }
        fmt.Println(count)
    }
}

func isPrime(n int) bool {
    isPrime := true
    for i := 2; i*i < n+1; i++ { // 2~√nまでの間で割れるものがなければ素数
        if n%i == 0 {
            isPrime = false
            return isPrime
        }
    }
    return isPrime
}

上のコードでのボトルネックは、明らかに2~√nまでの数字で割っている部分である。

このように素数判定をいくつもの数字でやる必要がある場合は「エラトステネスの篩」を使うと高速に素数判定ができる。

エラトステネスの篩

エラトステネスの篩は、配列にそのインデックスの数が素数かどうかを記録した表のこと。

配列の中身をtrueで初期化して、素数でない数字を記録していく。

素数でない数字はfor文で整数を回してその倍数を素数でない数字として記録する。

コードで表現すると以下のようになる。

// エラトステネスの篩
prime := make([]bool, 100001)     // 0 ~ 100000までの数字が素数かどうかを表で表す
for i := 2; i < len(prime); i++ { // 0, 1以外を素数として初期化
    prime[i] = true
}
for i := 2; i < len(prime); i++ {
    if !prime[i] { // 例えば4は既に2の倍数として素数じゃないと判定されているので考えなくてOK
        continue
    }
    for j := i * 2; j < len(prime); j += i { // iの倍数が素数でないことを記録
        prime[j] = false
    }
}

このエラトステネスの篩を使うことで高速に素数判定できる。

これを踏まえると次のようなコードでACできる。

package main

import "fmt"

func main() {
    var Q int
    fmt.Scan(&Q)
    l := make([]int, Q)
    r := make([]int, Q)
    for i := 0; i < Q; i++ {
        fmt.Scan(&l[i], &r[i])
    }

    // エラトステネスの篩
    prime := make([]bool, 100001)     // 0 ~ 100000までの数字が素数かどうかを表で表す
    for i := 2; i < len(prime); i++ { // 0, 1以外を素数として初期化
        prime[i] = true
    }
    for i := 2; i < len(prime); i++ {
        if !prime[i] { // 例えば4は既に2の倍数として素数じゃないと判定されているので考えなくてOK
            continue
        }
        for j := i * 2; j < len(prime); j += i { // iの倍数が素数でないことを記録
            prime[j] = false
        }
    }

    similar := make([]bool, 100001)
    for x := 1; x < len(similar); x += 2 { // 奇数だけを見ていく
        if prime[x] && prime[(x+1)/2] {
            similar[x] = true
        }
    }
    s := make([]int, len(similar)+1) // 累積和
    for i := 0; i < len(similar); i++ {
        if similar[i] {
            s[i+1] = s[i] + 1
        } else {
            s[i+1] = s[i]
        }
    }

    for i := 0; i < Q; i++ {
        fmt.Println(s[r[i]+1] - s[l[i]])
    }
}