nekolib/math/
prime_pi_.rs

1//! 素数の数え上げ。
2
3use super::linear_sieve;
4use super::super::traits::bisect;
5
6use std::ops::RangeFrom;
7
8use bisect::Bisect;
9use linear_sieve::LinearSieve;
10
11/// 素数の数え上げ。
12///
13/// $n$ 以下の素数の個数 $\\pi(n)$ を返す。
14///
15/// # Complexity
16/// $O(\\sqrt{n})$ space, $O(n^{2/3} / \\log(n)^{1/3})$ time.
17///
18/// # Examples
19/// ```
20/// use nekolib::math::prime_pi;
21///
22/// assert_eq!(prime_pi(10), 4);
23/// assert_eq!(prime_pi(100), 25);
24/// assert_eq!(prime_pi(1000), 168);
25/// assert_eq!(prime_pi(10000), 1229);
26/// assert_eq!(prime_pi(100_000_000_000), 4118054813);
27/// ```
28///
29/// # References
30/// - <https://rsk0315.github.io/slides/prime-counting.pdf>
31///
32/// # See also
33/// - <https://rsk0315.hatenablog.com/entry/2021/05/18/015511>
34/// - <https://judge.yosupo.jp/submission/7976>
35/// - <https://math314.hateblo.jp/entry/2016/06/05/004332>
36/// - <https://projecteuler.net/thread=10;page=5#111677>
37pub fn prime_pi(n: usize) -> usize {
38    let n_2 = (1_usize..).bisect(|i| i.pow(2) <= n) - 1;
39    let n_6 = (1_usize..).bisect(|i| i.pow(6) <= n) - 1;
40
41    let lg = 1.max(n.next_power_of_two().trailing_zeros() as usize);
42    let nlg_3 = (1_usize..).bisect(|i| i.pow(3) <= n * lg) - 1;
43    let n_lg2_3 = (1_usize..).bisect(|i| i.pow(3) * lg.pow(2) <= n) - 1;
44
45    let mut h = vec![0];
46    h.extend((1..=n_2).map(|i| n / i));
47    h.extend((1..n / n_2).rev());
48    let ns = h.clone();
49    for hi in &mut h[1..] {
50        *hi -= 1;
51    }
52
53    let primes: Vec<_> = LinearSieve::new(n_2).primes().collect();
54
55    let mut pi = 0;
56    while pi < primes.len() && primes[pi] <= n_6 {
57        let p = primes[pi];
58        let pp = p * p;
59        for (&n_i, i) in ns[1..].iter().take_while(|&&n_i| n_i >= pp).zip(1..) {
60            let ip = i * p;
61            h[i] -= h[if ip <= n_2 { ip } else { ns.len() - n_i / p }] - pi;
62        }
63        pi += 1;
64    }
65
66    let thresh = if nlg_3 <= n_2 { nlg_3 } else { ns.len() - n / nlg_3 };
67    let mut rsq = FenwickTree::new(ns.len() - thresh);
68    while pi < primes.len() && primes[pi] <= n_lg2_3 {
69        let p = primes[pi];
70        let pp = p * p;
71        for (&n_i, i) in
72            ns[1..].iter().take_while(|&&n_i| n_i >= pp).zip(1..=thresh)
73        {
74            let ip = i * p;
75            let index = if ip <= n_2 { ip } else { ns.len() - n_i / p };
76
77            let mut sum: usize = h[index];
78            if index > thresh {
79                sum -= rsq.sum(index - thresh..);
80            }
81            h[i] -= sum - pi;
82        }
83
84        let mut st = vec![(p, pi)];
85        while let Some((cur, i)) = st.pop() {
86            for (cur, j) in primes[i..]
87                .iter()
88                .map(|&p_j| cur * p_j)
89                .take_while(|&cur| cur < n / nlg_3)
90                .zip(i..)
91            {
92                let index = if cur <= n_2 { ns.len() - cur } else { n / cur };
93                if index > thresh {
94                    rsq.add(index - thresh, 1);
95                }
96                st.push((cur, j));
97            }
98        }
99
100        pi += 1;
101    }
102
103    let rsq = rsq.into_suffix_sum();
104    for i in 0..rsq.len() {
105        h[i + thresh] -= rsq[i];
106    }
107
108    while pi < primes.len() && primes[pi] <= n_2 {
109        let p = primes[pi];
110        let pp = p * p;
111        for (&n_i, i) in ns[1..].iter().take_while(|&&n_i| n_i >= pp).zip(1..) {
112            let ip = i * p;
113            h[i] -= h[if ip <= n_2 { ip } else { ns.len() - n_i / p }] - pi;
114        }
115        pi += 1;
116    }
117
118    h[1]
119}
120
121struct FenwickTree(Vec<usize>);
122
123impl FenwickTree {
124    pub fn new(n: usize) -> Self { Self(vec![0; n]) }
125    pub fn sum(&self, rf: RangeFrom<usize>) -> usize {
126        let mut i = rf.start;
127        let mut res = 0;
128        while i < self.0.len() {
129            res += self.0[i];
130            i += i & i.wrapping_neg();
131        }
132        res
133    }
134    pub fn add(&mut self, mut i: usize, d: usize) {
135        while i > 0 {
136            self.0[i] += d;
137            i -= i & i.wrapping_neg();
138        }
139    }
140    pub fn into_suffix_sum(self) -> Vec<usize> {
141        let mut res = self.0;
142        for i in (1..res.len()).rev() {
143            let j = i + (i & i.wrapping_neg());
144            if j < res.len() {
145                res[i] += res[j];
146            }
147        }
148        res
149    }
150}
nekolib/math/prime_pi_.rs

nekolib/math/
prime_pi_.rs
