えびちゃんの日記

えびちゃん(競プロ)の日記です。

n 番目の素数を求めるよ

軽めの記事です。

\(n\) 番目の素数を \(p_n\)、\(n\) 以下の素数を \(\pi(n)\) と書くとします。 このとき、\(p_n\) は \(\pi(p_n) \ge n\) かつ \(\pi(p_n-1) \lt n\) を満たします。 よって、これは二分探索で求められます。

上限を決めるのが大変そうですが、指数探索っぽくすればいいです。 指数探索というのは、上限を \(2, 4, 8, \dots\) と倍々に増やしていき、true/false の境界を跨いだらその範囲で二分探索を行う方法です。 今回は、自明に \(p_n \gt n\) なので、上限の初期値はそれに基づいて決めてよいです(別に \(2\) から始めてもよいです)。

fn main() {
    assert_eq!(nth_prime(3), 5);
    assert_eq!(nth_prime(4118054813), 99999999977);
}

fn nth_prime(n: usize) -> usize {
    let mut lo = n;
    let mut hi = 2 * n;
    let small = |i| prime_pi(i) < n;
    while small(hi) {
        hi *= 2;
    }
    while hi - lo > 1 {
        let mid = lo + (hi - lo) / 2;
        if small(mid) {
            lo = mid;
        } else {
            hi = mid;
        }
    }
    hi
}

ここで、prime_pi別の記事 で書いたやつを使うと、\(O(n^{3/4}/\log(n))\) time です。 また、\(p_n = O(n\log(n))\) が知られているので、全体で \(O( (n\log(n))^{3/4}/\log(n\log(n))\cdot \log(n\log(n)) = O( (n\log(n))^{3/4})\) time です。

お粗末さまでした。