Struct PrimeList

Source

pub struct PrimeList {
    bits: Vec<u64>,
    bit_len: usize,
    max_n: u32,
    prime_count: usize,
}

Fields§

§bits: Vec<u64>§bit_len: usize§max_n: u32§prime_count: usize

Implementations§

Source §

impl PrimeList

Source

pub fn new(max_n: u32) -> Self

Examples found in repository ?

crates/competitive/src/math/discrete_logarithm.rs (line 62)

60    fn new() -> Self {
61        Self {
62            primes: PrimeList::new(2),
63            br_primes: Default::default(),
64            ic: Default::default(),
65        }
66    }

More examples

Hide additional examples

crates/library_checker/src/number_theory/enumerate_primes.rs (line 8)

5pub fn enumerate_primes(reader: impl Read, writer: impl Write) {
6    prepare_io!(reader, writer);
7    sc!(n: u32, a, b);
8    let primes = PrimeList::new(n);
9    let iter = primes.primes().skip(b).step_by(a);
10    pp!(primes.len(), primes.len().saturating_sub(b).div_ceil(a); @it iter);
11}

Source

pub fn primes(&self) -> PrimeListIter<'_> ⓘ

Examples found in repository ?

crates/competitive/src/math/prime_list.rs (line 326)

322    pub fn trial_division(&self, n: u64) -> PrimeListTrialDivision<'_> {
323        let bound = u64::from(self.max_n).pow(2);
324        assert!(n <= bound, "expected `n={} <= {}`", n, bound);
325        PrimeListTrialDivision {
326            primes: self.primes(),
327            n,
328        }
329    }

More examples

Hide additional examples

crates/library_checker/src/number_theory/enumerate_primes.rs (line 9)

5pub fn enumerate_primes(reader: impl Read, writer: impl Write) {
6    prepare_io!(reader, writer);
7    sc!(n: u32, a, b);
8    let primes = PrimeList::new(n);
9    let iter = primes.primes().skip(b).step_by(a);
10    pp!(primes.len(), primes.len().saturating_sub(b).div_ceil(a); @it iter);
11}

Source

pub fn len(&self) -> usize

Examples found in repository ?

crates/library_checker/src/number_theory/enumerate_primes.rs (line 10)

5pub fn enumerate_primes(reader: impl Read, writer: impl Write) {
6    prepare_io!(reader, writer);
7    sc!(n: u32, a, b);
8    let primes = PrimeList::new(n);
9    let iter = primes.primes().skip(b).step_by(a);
10    pp!(primes.len(), primes.len().saturating_sub(b).div_ceil(a); @it iter);
11}

Source

pub fn is_empty(&self) -> bool

Source

pub fn primes_lte(&self, n: u32) -> PrimeListIter<'_> ⓘ

Examples found in repository ?

crates/competitive/src/math/prime_list.rs (line 277)

276    pub fn primes(&self) -> PrimeListIter<'_> {
277        self.primes_lte(self.max_n)
278    }

More examples

Hide additional examples

crates/competitive/src/math/lcm_convolve.rs (line 16)

13    pub fn zeta_transform(f: &mut [M::T]) {
14        let n = f.len().saturating_sub(1) as u32;
15        with_prime_list(n, |pl| {
16            for p in pl.primes_lte(n) {
17                for (i, j) in (0..f.len()).step_by(p as _).enumerate() {
18                    f[j] = M::operate(&f[j], &f[i]);
19                }
20            }
21        })
22    }
23}
24
25impl<G> LcmConvolve<G>
26where
27    G: Group,
28{
29    /// $$f(m) = \sum_{n \mid m}h(n)$$
30    pub fn mobius_transform(f: &mut [G::T]) {
31        let n = f.len().saturating_sub(1) as u32;
32        with_prime_list(n, |pl| {
33            for p in pl.primes_lte(n) {
34                for (i, j) in (0..f.len()).step_by(p as _).enumerate().rev() {
35                    f[j] = G::rinv_operate(&f[j], &f[i]);
36                }
37            }
38        })
39    }

crates/competitive/src/math/gcd_convolve.rs (line 16)

13    pub fn zeta_transform(f: &mut [M::T]) {
14        let n = f.len().saturating_sub(1) as u32;
15        with_prime_list(n, |pl| {
16            for p in pl.primes_lte(n) {
17                for (i, j) in (0..f.len()).step_by(p as _).enumerate().rev() {
18                    f[i] = M::operate(&f[i], &f[j]);
19                }
20            }
21        })
22    }
23}
24
25impl<G> GcdConvolve<G>
26where
27    G: Group,
28{
29    /// $$f(m) = \sum_{n \mid m}h(n)$$
30    pub fn mobius_transform(f: &mut [G::T]) {
31        let n = f.len().saturating_sub(1) as u32;
32        with_prime_list(n, |pl| {
33            for p in pl.primes_lte(n) {
34                for (i, j) in (0..f.len()).step_by(p as _).enumerate() {
35                    f[i] = G::rinv_operate(&f[i], &f[j]);
36                }
37            }
38        })
39    }

crates/competitive/src/math/discrete_logarithm.rs (line 71)

67    fn discrete_logarithm(&mut self, a: u64, b: u64, p: u64) -> Option<(u64, u64)> {
68        let lim = ((((p as f64).log2() * (p as f64).log2().log2()).sqrt() / 2.0 + 1.).exp2() * 0.9)
69            as u32;
70        self.primes.reserve(lim);
71        let prime_count = self.primes.primes_lte(lim).count();
72        self.br_primes.extend(
73            self.primes
74                .primes_lte(lim)
75                .skip(self.br_primes.len())
76                .map(|p| BarrettReduction::<u64>::new(p.into())),
77        );
78        let br_primes = &self.br_primes[..prime_count];
79        self.ic
80            .entry(p)
81            .or_insert_with(|| IndexCalculusWithPrimitiveRoot::new(p, br_primes))
82            .discrete_logarithm(a, b, br_primes)
83    }

crates/competitive/src/math/quotient_array.rs (line 67)

61    pub fn lucy_dp<G>(mut self, mut mul_p: impl FnMut(T, u64) -> T) -> Self
62    where
63        G: Group<T = T>,
64    {
65        let max_n = self.isqrtn as u32;
66        with_prime_list(max_n, |pl| {
67            for p in pl.primes_lte(max_n) {
68                let p = u64::from(p);
69                let k = self.quotient_index(p - 1);
70                let p2 = p * p;
71                for (i, q) in Self::index_iter(self.n, self.isqrtn).enumerate() {
72                    if q < p2 {
73                        break;
74                    }
75                    let diff = mul_p(G::rinv_operate(&self[q / p], &self.data[k]), p);
76                    G::rinv_operate_assign(&mut self.data[i], &diff);
77                }
78            }
79        });
80        self
81    }
82
83    /// convert $\sum_{i\leq n, i\text{ is prime}} f(i)$ to $\sum_{i\leq n} f(i)$
84    pub fn min_25_sieve<R>(&self, mut f: impl FnMut(u64, u32) -> T) -> Self
85    where
86        T: Clone + One,
87        R: Ring<T = T, Additive: Invertible>,
88    {
89        let mut dp = self.clone();
90        let max_n = self.isqrtn as u32;
91        with_prime_list(max_n, |pl| {
92            for p in pl.primes_lte(max_n).rev() {
93                let p = u64::from(p);
94                let k = self.quotient_index(p);
95                for (i, q) in Self::index_iter(self.n, self.isqrtn).enumerate() {
96                    let mut pc = p;
97                    if pc * p > q {
98                        break;
99                    }
100                    let mut c = 1;
101                    while q / p >= pc {
102                        let x = R::mul(&f(p, c), &(R::sub(&dp[q / pc], &self.data[k])));
103                        let x = R::add(&x, &f(p, c + 1));
104                        dp.data[i] = R::add(&dp.data[i], &x);
105                        c += 1;
106                        pc *= p;
107                    }
108                }
109            }
110        });
111        for x in &mut dp.data {
112            *x = R::add(x, &T::one());
113        }
114        dp
115    }

Source

pub fn is_prime(&self, n: u32) -> bool

Source

pub fn trial_division(&self, n: u64) -> PrimeListTrialDivision<'_> ⓘ

Examples found in repository ?

crates/competitive/src/math/prime_list.rs (line 331)

330    pub fn prime_factors(&self, n: u64) -> Vec<(u64, u32)> {
331        self.trial_division(n).collect()
332    }
333    pub fn count_divisors(&self, n: u64) -> u64 {
334        let mut divisor_cnt = 1u64;
335        for (_, cnt) in self.trial_division(n) {
336            divisor_cnt *= cnt as u64 + 1;
337        }
338        divisor_cnt
339    }
340    pub fn divisors(&self, n: u64) -> Vec<u64> {
341        let mut d = vec![1u64];
342        for (p, c) in self.trial_division(n) {
343            let k = d.len();
344            let mut acc = 1;
345            for _ in 0..c {
346                acc *= p;
347                for i in 0..k {
348                    d.push(d[i] * acc);
349                }
350            }
351        }
352        d.sort_unstable();
353        d
354    }

Source

pub fn reserve(&mut self, max_n: u32)

Extends the prime list up to max_n.

Examples found in repository ?

crates/competitive/src/math/prime_list.rs (line 273)

271    pub fn new(max_n: u32) -> Self {
272        let mut self_: Self = Default::default();
273        self_.reserve(max_n);
274        self_
275    }
276    pub fn primes(&self) -> PrimeListIter<'_> {
277        self.primes_lte(self.max_n)
278    }
279    pub fn len(&self) -> usize {
280        self.prime_count
281    }
282    pub fn is_empty(&self) -> bool {
283        self.prime_count == 0
284    }
285    pub fn primes_lte(&self, n: u32) -> PrimeListIter<'_> {
286        assert!(n <= self.max_n, "expected `n={} <= {}`", n, self.max_n);
287        let bit_len = to_ordinal(n.saturating_add(1)) as usize;
288        let words = &self.bits[..self.bits.len().min(bit_len.div_ceil(64))];
289        let last_mask = if bit_len.is_multiple_of(64) {
290            !0
291        } else {
292            (1 << (bit_len % 64)) - 1
293        };
294        let (front_word, middle_words, back_word) = match words {
295            [] => (0, words, 0),
296            [word] => (*word & last_mask, &words[1..], 0),
297            [front, middle @ .., back] => (*front, middle, *back & last_mask),
298        };
299        let back_word_ordinal = words.len().saturating_sub(1) as u32 * 64;
300        PrimeListIter {
301            wheel_indices: 0..WHEEL_PRIMES.partition_point(|&p| p <= n) as u8,
302            middle_words,
303            front_word_base: 0,
304            front_word_phase: 0,
305            front_word,
306            back_word_base: back_word_ordinal / COPRIME as u32 * PERIOD,
307            back_word_phase: (back_word_ordinal % COPRIME as u32) as u8,
308            back_word,
309        }
310    }
311    pub fn is_prime(&self, n: u32) -> bool {
312        assert!(n <= self.max_n, "expected `n={} <= {}`", n, self.max_n);
313        if WHEEL_PRIMES.contains(&n) {
314            true
315        } else if !coprime_to_wheel(n) {
316            false
317        } else {
318            let ordinal = to_ordinal(n) as usize;
319            ordinal < self.bit_len && self.bits[ordinal / 64] >> (ordinal % 64) & 1 != 0
320        }
321    }
322    pub fn trial_division(&self, n: u64) -> PrimeListTrialDivision<'_> {
323        let bound = u64::from(self.max_n).pow(2);
324        assert!(n <= bound, "expected `n={} <= {}`", n, bound);
325        PrimeListTrialDivision {
326            primes: self.primes(),
327            n,
328        }
329    }
330    pub fn prime_factors(&self, n: u64) -> Vec<(u64, u32)> {
331        self.trial_division(n).collect()
332    }
333    pub fn count_divisors(&self, n: u64) -> u64 {
334        let mut divisor_cnt = 1u64;
335        for (_, cnt) in self.trial_division(n) {
336            divisor_cnt *= cnt as u64 + 1;
337        }
338        divisor_cnt
339    }
340    pub fn divisors(&self, n: u64) -> Vec<u64> {
341        let mut d = vec![1u64];
342        for (p, c) in self.trial_division(n) {
343            let k = d.len();
344            let mut acc = 1;
345            for _ in 0..c {
346                acc *= p;
347                for i in 0..k {
348                    d.push(d[i] * acc);
349                }
350            }
351        }
352        d.sort_unstable();
353        d
354    }
355    /// Extends the prime list up to `max_n`.
356    pub fn reserve(&mut self, max_n: u32) {
357        if max_n <= self.max_n || max_n < 2 {
358            return;
359        }
360        let limit = max_n.saturating_add(1);
361        self.bit_len = to_ordinal(limit) as usize;
362        if limit <= SQRT_THRESHOLD {
363            self.bits = SQRT_BITS[..self.bit_len.div_ceil(64)].to_vec();
364        } else {
365            self.bits = vec![!0; self.bit_len.div_ceil(64)];
366            let (wheels, medium_primes_begin) = make_wheels();
367            const DENSE_BLOCK: u32 = 1 << 25;
368            for start in (0..limit).step_by(DENSE_BLOCK as usize) {
369                let right = start.saturating_add(DENSE_BLOCK).min(limit);
370                for wheel in &wheels {
371                    let left = start / wheel.product * wheel.product;
372                    sieve_dense(&mut self.bits, to_ordinal(left), to_ordinal(right), wheel);
373                }
374            }
375
376            let steps = ordinal_steps();
377            let mut positions: Vec<_> = SQRT_PRIMES.iter().map(|&p| to_ordinal(p * p)).collect();
378            let mut states: Vec<_> = SQRT_PRIMES
379                .iter()
380                .map(|&p| STATES[(p % PERIOD) as usize])
381                .collect();
382            const SPARSE_BLOCK: u32 = 1 << 22;
383            for start in (0..limit).step_by(SPARSE_BLOCK as usize) {
384                let right = to_ordinal(start.saturating_add(SPARSE_BLOCK).min(limit));
385                for i in medium_primes_begin..SQRT_PRIME_COUNT {
386                    (positions[i], states[i]) =
387                        sieve_sparse(&mut self.bits, positions[i], right, i, states[i], &steps);
388                }
389            }
390            for (value, &sqrt_bits) in self.bits.iter_mut().zip(&SQRT_BITS) {
391                *value = sqrt_bits;
392            }
393        }
394
395        self.prime_count = WHEEL_PRIMES.partition_point(|&p| p <= max_n);
396        if let Some((&last, rest)) = self.bits.split_last() {
397            self.prime_count += rest
398                .iter()
399                .map(|word| word.count_ones() as usize)
400                .sum::<usize>();
401            let last_mask = if self.bit_len.is_multiple_of(64) {
402                !0
403            } else {
404                (1 << (self.bit_len % 64)) - 1
405            };
406            self.prime_count += (last & last_mask).count_ones() as usize;
407        }
408        self.max_n = max_n;
409    }
410}
411
412#[derive(Clone, Debug)]
413pub struct PrimeListIter<'a> {
414    wheel_indices: Range<u8>,
415    middle_words: &'a [u64],
416    front_word_base: u32,
417    front_word_phase: u8,
418    front_word: u64,
419    back_word_base: u32,
420    back_word_phase: u8,
421    back_word: u64,
422}
423
424impl PrimeListIter<'_> {
425    #[inline(always)]
426    fn load_front_word(&mut self) -> bool {
427        if let Some((&word, words)) = self.middle_words.split_first() {
428            self.middle_words = words;
429            if self.front_word_phase == 32 {
430                self.front_word_base += PERIOD * 2;
431                self.front_word_phase = 0;
432            } else {
433                self.front_word_base += PERIOD;
434                self.front_word_phase += 16;
435            }
436            self.front_word = word;
437            true
438        } else if self.back_word != 0 {
439            self.front_word_base = self.back_word_base;
440            self.front_word_phase = self.back_word_phase;
441            self.front_word = replace(&mut self.back_word, 0);
442            true
443        } else {
444            false
445        }
446    }
447
448    #[inline(always)]
449    fn load_back_word(&mut self) -> bool {
450        if let Some((&word, words)) = self.middle_words.split_last() {
451            self.middle_words = words;
452            if self.back_word_phase == 0 {
453                self.back_word_base -= PERIOD * 2;
454                self.back_word_phase = 32;
455            } else {
456                self.back_word_base -= PERIOD;
457                self.back_word_phase -= 16;
458            }
459            self.back_word = word;
460            true
461        } else if self.front_word != 0 {
462            self.back_word_base = self.front_word_base;
463            self.back_word_phase = self.front_word_phase;
464            self.back_word = replace(&mut self.front_word, 0);
465            true
466        } else {
467            false
468        }
469    }
470}
471
472impl Iterator for PrimeListIter<'_> {
473    type Item = u32;
474
475    fn next(&mut self) -> Option<Self::Item> {
476        if let Some(index) = self.wheel_indices.next() {
477            return Some(WHEEL_PRIMES[index as usize]);
478        }
479        loop {
480            if self.front_word != 0 {
481                let bit = self.front_word.trailing_zeros();
482                self.front_word &= self.front_word - 1;
483                return Some(
484                    self.front_word_base
485                        + ORDINAL_TO_VALUE[(self.front_word_phase + bit as u8) as usize] as u32,
486                );
487            }
488            if !self.load_front_word() {
489                return None;
490            }
491        }
492    }
493
494    fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
495        if n == 0 {
496            return self.next();
497        }
498        for index in self.wheel_indices.by_ref() {
499            if n == 0 {
500                return Some(WHEEL_PRIMES[index as usize]);
501            }
502            n -= 1;
503        }
504        loop {
505            let count = self.front_word.count_ones() as usize;
506            if n < count {
507                for _ in 0..n {
508                    self.front_word &= self.front_word - 1;
509                }
510                return self.next();
511            }
512            n -= count;
513            self.front_word = 0;
514            if !self.load_front_word() {
515                return None;
516            }
517        }
518    }
519}
520
521impl DoubleEndedIterator for PrimeListIter<'_> {
522    fn next_back(&mut self) -> Option<Self::Item> {
523        loop {
524            if self.back_word != 0 {
525                let bit = 63 - self.back_word.leading_zeros();
526                self.back_word -= 1 << bit;
527                return Some(
528                    self.back_word_base
529                        + ORDINAL_TO_VALUE[(self.back_word_phase + bit as u8) as usize] as u32,
530                );
531            }
532            if !self.load_back_word() {
533                return self
534                    .wheel_indices
535                    .next_back()
536                    .map(|index| WHEEL_PRIMES[index as usize]);
537            }
538        }
539    }
540
541    fn nth_back(&mut self, mut n: usize) -> Option<Self::Item> {
542        if n == 0 {
543            return self.next_back();
544        }
545        loop {
546            let count = self.back_word.count_ones() as usize;
547            if n < count {
548                for _ in 0..n {
549                    let bit = 63 - self.back_word.leading_zeros();
550                    self.back_word -= 1 << bit;
551                }
552                return self.next_back();
553            }
554            n -= count;
555            self.back_word = 0;
556            if !self.load_back_word() {
557                for index in self.wheel_indices.by_ref().rev() {
558                    if n == 0 {
559                        return Some(WHEEL_PRIMES[index as usize]);
560                    }
561                    n -= 1;
562                }
563                return None;
564            }
565        }
566    }
567}
568
569#[derive(Debug, Clone)]
570pub struct PrimeListTrialDivision<'p> {
571    primes: PrimeListIter<'p>,
572    n: u64,
573}
574impl Iterator for PrimeListTrialDivision<'_> {
575    type Item = (u64, u32);
576    fn next(&mut self) -> Option<Self::Item> {
577        if self.n <= 1 {
578            return None;
579        }
580        for p in self.primes.by_ref() {
581            let p = u64::from(p);
582            if p * p > self.n {
583                break;
584            }
585            if self.n.is_multiple_of(p) {
586                let mut cnt = 1u32;
587                self.n /= p;
588                while self.n.is_multiple_of(p) {
589                    cnt += 1;
590                    self.n /= p;
591                }
592                return Some((p, cnt));
593            }
594        }
595        if self.n > 1 {
596            return Some((replace(&mut self.n, 1), 1));
597        }
598        None
599    }
600}
601
602pub fn with_prime_list<F>(max_n: u32, f: F)
603where
604    F: FnOnce(&PrimeList),
605{
606    thread_local!(static PRIME_LIST: UnsafeCell<PrimeList> = Default::default());
607    PRIME_LIST.with(|cell| {
608        unsafe {
609            let pl = &mut *cell.get();
610            pl.reserve(max_n);
611            f(pl);
612        };
613    });
614}

More examples

Hide additional examples

crates/competitive/src/math/discrete_logarithm.rs (line 70)

67    fn discrete_logarithm(&mut self, a: u64, b: u64, p: u64) -> Option<(u64, u64)> {
68        let lim = ((((p as f64).log2() * (p as f64).log2().log2()).sqrt() / 2.0 + 1.).exp2() * 0.9)
69            as u32;
70        self.primes.reserve(lim);
71        let prime_count = self.primes.primes_lte(lim).count();
72        self.br_primes.extend(
73            self.primes
74                .primes_lte(lim)
75                .skip(self.br_primes.len())
76                .map(|p| BarrettReduction::<u64>::new(p.into())),
77        );
78        let br_primes = &self.br_primes[..prime_count];
79        self.ic
80            .entry(p)
81            .or_insert_with(|| IndexCalculusWithPrimitiveRoot::new(p, br_primes))
82            .discrete_logarithm(a, b, br_primes)
83    }