素数列生成

素数列の実装方法を考えてみたのでメモ。

基本的にはエラトステネスの篩と同じようなことを行っているはずだが、一度に保持する篩の範囲が(おそらく)必要最小限となっている。
(一つの既知の素数につき、ハッシュテーブル内の一つの枠しか消費しない)

// file: prime.rs
//
// $ rustc --version
// rustc 1.0.0 (a59de37e9 2015-05-13) (built 2015-05-14)

use std::collections::HashMap;

pub struct Prime {
    curr: u64, // 現在の数値 (素数とは限らない)
    sieve: HashMap<u64, u64>, // 合成数 => 素数: サイズは`curr未満の素数の数`に常に等しい
}

impl Prime {
    pub fn is_prime(&mut self, n: u64) -> bool {
        match self.sieve.remove(&n) { // 一度探索した数はsieveから除去する(容量節約)
            None        => {self.mark_composite(2, n);               true},  // 素数
            Some(prime) => {self.mark_composite(n/prime + 1, prime); false}, // 合成数
        }
    }
    fn mark_composite(&mut self, times: u64, prime: u64) {
        (times..).find(|i| !self.sieve.contains_key(&(i*prime)) ) // 未登録の合成数を探す
                 .and_then(|i| self.sieve.insert(i*prime, prime) );
    }
}

impl Iterator for Prime {
    type Item = u64;
    fn next(&mut self) -> Option<u64> {
        let prime = (self.curr..).find(|n| self.is_prime(*n) ).unwrap(); // curr以降の最小の素数を検索
        self.curr = prime + 1;
        Some(prime)
    }
}

pub fn primes() -> Prime {
    Prime{curr: 2, sieve: HashMap::new()}
}

fn main() {
    use std::env;
    let n: usize = env::args().nth(1).unwrap().parse().unwrap();
    println!("{}-th prime: {}", n, primes().nth(n-1).unwrap());
}

実行例:

# CPU: Intel(R) Core(TM) i7-5600U CPU @ 2.60GHz

# コンパイル
$ rustc -O hoge.rs

# 100万番の素数の検索
$ time ./hoge 1000000
1000000-th prime: 15485863

real    0m10.086s  # 約10秒掛かった
user    0m10.060s
sys 0m0.028s