competitive/math/
prime_list.rs

1use std::{cell::UnsafeCell, mem::replace, ops::Range};
2
3const WHEEL_PRIMES: [u32; 4] = [2, 3, 5, 7];
4const PERIOD: u32 = 210;
5const COPRIME: usize = 48;
6const SQRT_THRESHOLD: u32 = 1 << 16;
7
8const fn coprime_to_wheel(x: u32) -> bool {
9    !x.is_multiple_of(2) && !x.is_multiple_of(3) && !x.is_multiple_of(5) && !x.is_multiple_of(7)
10}
11
12const fn residues() -> [u8; COPRIME] {
13    let mut result = [0; COPRIME];
14    let mut i = 1;
15    let mut j = 0;
16    while i < PERIOD {
17        if coprime_to_wheel(i) {
18            result[j] = i as u8;
19            j += 1;
20        }
21        i += 2;
22    }
23    result
24}
25
26const RESIDUES: [u8; COPRIME] = residues();
27
28const fn states() -> [u8; PERIOD as usize] {
29    let mut result = [0; PERIOD as usize];
30    let mut i = 0;
31    let mut j = 0;
32    while i < PERIOD {
33        result[i as usize] = j;
34        if coprime_to_wheel(i) {
35            j += 1;
36        }
37        i += 1;
38    }
39    result
40}
41
42const STATES: [u8; PERIOD as usize] = states();
43
44const fn additions() -> [u8; PERIOD as usize] {
45    let mut result = [0; PERIOD as usize];
46    let mut i = 0;
47    while i < PERIOD {
48        let mut add = 1;
49        while !coprime_to_wheel(i + add) {
50            add += 1;
51        }
52        result[i as usize] = add as u8;
53        i += 1;
54    }
55    result
56}
57
58const ADDITIONS: [u8; PERIOD as usize] = additions();
59
60const fn gaps() -> [u8; COPRIME] {
61    let mut result = [0; COPRIME];
62    let mut i = 0;
63    while i < COPRIME {
64        result[i] = ADDITIONS[RESIDUES[i] as usize];
65        i += 1;
66    }
67    result
68}
69
70const GAPS: [u8; COPRIME] = gaps();
71
72const fn to_ordinal(x: u32) -> u32 {
73    x / PERIOD * COPRIME as u32 + STATES[(x % PERIOD) as usize] as u32
74}
75
76const fn to_value(x: u32) -> u32 {
77    x / COPRIME as u32 * PERIOD + RESIDUES[x as usize % COPRIME] as u32
78}
79
80const fn ordinal_to_value() -> [u16; 256] {
81    let mut result = [0; 256];
82    let mut i = 0;
83    while i < result.len() {
84        result[i] = to_value(i as u32) as u16;
85        i += 1;
86    }
87    result
88}
89
90const ORDINAL_TO_VALUE: [u16; 256] = ordinal_to_value();
91
92const fn sqrt_bits() -> [u64; SQRT_THRESHOLD as usize / 128] {
93    let mut result = [!0; SQRT_THRESHOLD as usize / 128];
94    let ordinal = to_ordinal(1) as usize;
95    result[ordinal / 64] &= !(1 << (ordinal % 64));
96    let mut i = RESIDUES[1] as u32;
97    while to_ordinal(i * i) < SQRT_THRESHOLD / 2 {
98        let ordinal = to_ordinal(i) as usize;
99        if result[ordinal / 64] >> (ordinal % 64) & 1 != 0 {
100            let mut k = i;
101            while to_ordinal(i * k) < SQRT_THRESHOLD / 2 {
102                let ordinal = to_ordinal(i * k) as usize;
103                result[ordinal / 64] &= !(1 << (ordinal % 64));
104                k += ADDITIONS[(k % PERIOD) as usize] as u32;
105            }
106        }
107        i += ADDITIONS[(i % PERIOD) as usize] as u32;
108    }
109    result
110}
111
112const SQRT_BITS: [u64; SQRT_THRESHOLD as usize / 128] = sqrt_bits();
113
114const fn count_sqrt_primes() -> usize {
115    let mut result = 0;
116    let mut i = RESIDUES[1] as u32;
117    while i < SQRT_THRESHOLD {
118        let ordinal = to_ordinal(i) as usize;
119        result += (SQRT_BITS[ordinal / 64] >> (ordinal % 64) & 1) as usize;
120        i += ADDITIONS[(i % PERIOD) as usize] as u32;
121    }
122    result
123}
124
125const SQRT_PRIME_COUNT: usize = count_sqrt_primes();
126
127const fn sqrt_primes() -> [u32; SQRT_PRIME_COUNT] {
128    let mut result = [0; SQRT_PRIME_COUNT];
129    let mut i = RESIDUES[1] as u32;
130    let mut j = 0;
131    while i < SQRT_THRESHOLD {
132        let ordinal = to_ordinal(i) as usize;
133        if SQRT_BITS[ordinal / 64] >> (ordinal % 64) & 1 != 0 {
134            result[j] = i;
135            j += 1;
136        }
137        i += ADDITIONS[(i % PERIOD) as usize] as u32;
138    }
139    result
140}
141
142static SQRT_PRIMES: [u32; SQRT_PRIME_COUNT] = sqrt_primes();
143
144struct Wheel {
145    mask: Vec<u64>,
146    product: u32,
147}
148
149impl Wheel {
150    fn new(primes: &[u32], product: u32) -> Self {
151        let mut mask = vec![!0; to_ordinal(product) as usize / 64];
152        for &p in primes {
153            let mut k = 1;
154            while p * k < product {
155                let ordinal = to_ordinal(p * k) as usize;
156                mask[ordinal / 64] &= !(1 << (ordinal % 64));
157                k += ADDITIONS[(k % PERIOD) as usize] as u32;
158            }
159        }
160        Self { mask, product }
161    }
162}
163
164fn make_wheels() -> (Vec<Wheel>, usize) {
165    const MAX_WHEEL_SIZE: u32 = 1 << 20;
166    const BASE: u32 = (PERIOD * 64) >> (WHEEL_PRIMES.len() - 2);
167    let mut product = BASE;
168    let mut current = vec![];
169    let mut wheels = vec![];
170    for (i, &p) in SQRT_PRIMES.iter().enumerate() {
171        if product * p > MAX_WHEEL_SIZE {
172            wheels.push(Wheel::new(&current, product));
173            current.clear();
174            current.push(p);
175            product = BASE * p;
176            if product > MAX_WHEEL_SIZE {
177                return (wheels, i);
178            }
179        } else {
180            current.push(p);
181            product *= p;
182        }
183    }
184    unreachable!()
185}
186
187fn sieve_dense(bits: &mut [u64], l: u32, r: u32, wheel: &Wheel) {
188    let mut left = l as usize / 64;
189    let right = (r as usize).div_ceil(64);
190    while left + wheel.mask.len() <= right {
191        for (value, &mask) in bits[left..left + wheel.mask.len()]
192            .iter_mut()
193            .zip(&wheel.mask)
194        {
195            *value &= mask;
196        }
197        left += wheel.mask.len();
198    }
199    for (value, &mask) in bits[left..right].iter_mut().zip(&wheel.mask) {
200        *value &= mask;
201    }
202}
203
204fn ordinal_steps() -> Vec<[u32; COPRIME * 2]> {
205    SQRT_PRIMES
206        .iter()
207        .map(|&p| {
208            let mut result = [0; COPRIME * 2];
209            let mut last = to_ordinal(p);
210            for i in 0..COPRIME {
211                let next = to_ordinal(p * (RESIDUES[i] as u32 + GAPS[i] as u32));
212                result[i] = next - last;
213                result[i + COPRIME] = next - last;
214                last = next;
215            }
216            result
217        })
218        .collect()
219}
220
221fn sieve_sparse(
222    bits: &mut [u64],
223    mut left: u32,
224    right: u32,
225    prime_index: usize,
226    mut state: u8,
227    steps: &[[u32; COPRIME * 2]],
228) -> (u32, u8) {
229    let p = SQRT_PRIMES[prime_index];
230    while left + p * COPRIME as u32 <= right {
231        for _ in 0..COPRIME {
232            let ordinal = left as usize;
233            bits[ordinal / 64] &= !(1 << (ordinal % 64));
234            left += steps[prime_index][state as usize];
235            state += 1;
236        }
237        state -= COPRIME as u8;
238    }
239    while left < right {
240        let ordinal = left as usize;
241        bits[ordinal / 64] &= !(1 << (ordinal % 64));
242        left += steps[prime_index][state as usize];
243        state += 1;
244    }
245    if state >= COPRIME as u8 {
246        state -= COPRIME as u8;
247    }
248    (left, state)
249}
250
251#[derive(Debug, Clone)]
252pub struct PrimeList {
253    bits: Vec<u64>,
254    bit_len: usize,
255    max_n: u32,
256    prime_count: usize,
257}
258
259impl Default for PrimeList {
260    fn default() -> Self {
261        Self {
262            bits: vec![],
263            bit_len: 0,
264            max_n: 1,
265            prime_count: 0,
266        }
267    }
268}
269
270impl PrimeList {
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}
615
616#[cfg(test)]
617mod tests {
618    use super::*;
619    use crate::math::prime_factors;
620    use crate::tools::{
621        Xorshift,
622        testutil::{exhaustive_sequences, sample_usize, structured_sequences},
623    };
624    use std::collections::VecDeque;
625
626    fn primes(n: usize) -> Vec<usize> {
627        if n < 2 {
628            return vec![];
629        }
630        let mut res = vec![2];
631        let sqrt_n = (n as f32).sqrt() as usize | 1;
632        let mut seive = vec![true; n / 2];
633        for i in (3..=sqrt_n).step_by(2) {
634            if seive[i / 2 - 1] {
635                res.push(i);
636                for j in (i * i..=n).step_by(i * 2) {
637                    seive[j / 2 - 1] = false;
638                }
639            }
640        }
641        for i in (std::cmp::max(3, sqrt_n + 2)..=n).step_by(2) {
642            if seive[i / 2 - 1] {
643                res.push(i);
644            }
645        }
646        res
647    }
648
649    pub fn divisors(n: u64) -> Vec<u64> {
650        let mut res = vec![];
651        for i in 1..(n as f32).sqrt() as u64 + 1 {
652            if n.is_multiple_of(i) {
653                res.push(i);
654                if i * i != n {
655                    res.push(n / i);
656                }
657            }
658        }
659        res.sort_unstable();
660        res
661    }
662
663    #[test]
664    fn test_prime_list() {
665        let mut rng = Xorshift::default();
666
667        for n in (0..1000).chain(rng.random_iter(0..=20000).take(100)) {
668            let pl = PrimeList::new(n);
669            let ps: Vec<_> = primes(n as _).into_iter().map(|p| p as u32).collect();
670            assert_eq!(pl.len(), ps.len());
671            assert_eq!(pl.primes().collect::<Vec<_>>(), ps);
672        }
673
674        for _ in 0..100 {
675            let b = rng.randf() * 0.0001;
676            let mut pl = PrimeList::new(0);
677            for n in (0..20_000).filter(|_| rng.gen_bool(b)) {
678                pl.reserve(n);
679                let ps: Vec<_> = primes(n as _).into_iter().map(|p| p as u32).collect();
680                assert_eq!(pl.len(), ps.len());
681                assert_eq!(pl.primes().collect::<Vec<_>>(), ps);
682            }
683        }
684
685        let pl = PrimeList::new(100_000);
686        for n in (0..1000).chain(rng.random_iter(0..=1_000_000_000).take(100)) {
687            assert_eq!(prime_factors(n), pl.prime_factors(n));
688        }
689    }
690
691    #[test]
692    fn test_primes() {
693        let mut rng = Xorshift::default();
694        let bounds = sample_usize(&mut rng, 128, 0..=100_000, 300);
695        for n in bounds {
696            let pl = PrimeList::new(n as u32);
697            let expected: Vec<_> = primes(n).into_iter().map(|p| p as u32).collect();
698            assert_eq!(pl.primes().collect::<Vec<_>>(), expected);
699            for limit in (0..=n.min(128)).chain([n / 2, n.saturating_sub(1), n]) {
700                let bounded: Vec<_> = expected
701                    .iter()
702                    .copied()
703                    .take_while(|&p| p <= limit as u32)
704                    .collect();
705                assert_eq!(pl.primes_lte(limit as u32).collect::<Vec<_>>(), bounded);
706                assert_eq!(
707                    pl.primes_lte(limit as u32).rev().collect::<Vec<_>>(),
708                    bounded.into_iter().rev().collect::<Vec<_>>()
709                );
710            }
711            for i in 0..=n.min(2000) {
712                assert_eq!(
713                    pl.is_prime(i as u32),
714                    expected.binary_search(&(i as u32)).is_ok()
715                );
716            }
717        }
718        // Exhaust all query limits on a larger, fixed sieve as well.
719        let pl = PrimeList::new(10_000);
720        let expected = primes(10_000);
721        for limit in 0..=10_000 {
722            let bounded: Vec<_> = expected
723                .iter()
724                .copied()
725                .take_while(|&p| p <= limit)
726                .map(|p| p as u32)
727                .collect();
728            assert_eq!(pl.primes_lte(limit as u32).collect::<Vec<_>>(), bounded);
729            assert_eq!(
730                pl.primes_lte(limit as u32).rev().collect::<Vec<_>>(),
731                bounded.into_iter().rev().collect::<Vec<_>>()
732            );
733        }
734    }
735
736    #[test]
737    fn test_prime_iterator() {
738        let mut rng = Xorshift::default();
739        for n in sample_usize(&mut rng, 128, 0..=10_000, 300) {
740            let pl = PrimeList::new(n as u32);
741            let expected: Vec<_> = primes(n).into_iter().map(|p| p as u32).collect();
742            for skip in (0..=expected.len().min(16)).chain([expected.len(), expected.len() + 1]) {
743                for step in (1..=16).chain([expected.len() + 1]) {
744                    assert_eq!(
745                        pl.primes().skip(skip).step_by(step).collect::<Vec<_>>(),
746                        expected
747                            .iter()
748                            .copied()
749                            .skip(skip)
750                            .step_by(step)
751                            .collect::<Vec<_>>()
752                    );
753                    assert_eq!(
754                        pl.primes()
755                            .rev()
756                            .skip(skip)
757                            .step_by(step)
758                            .collect::<Vec<_>>(),
759                        expected
760                            .iter()
761                            .rev()
762                            .copied()
763                            .skip(skip)
764                            .step_by(step)
765                            .collect::<Vec<_>>()
766                    );
767                }
768            }
769            let directions: Vec<Vec<_>> = if n <= 20 {
770                exhaustive_sequences(0..2, expected.len()..=expected.len()).collect()
771            } else {
772                structured_sequences(&mut rng, 0..2, [expected.len()]).collect()
773            };
774            for directions in directions {
775                let mut queue = VecDeque::from(expected.clone());
776                let mut iter = pl.primes();
777                for back in directions.into_iter().map(|i| i != 0) {
778                    if back {
779                        assert_eq!(iter.next_back(), queue.pop_back());
780                    } else {
781                        assert_eq!(iter.next(), queue.pop_front());
782                    }
783                }
784                assert!(queue.is_empty());
785                assert_eq!(iter.next(), None);
786                assert_eq!(iter.next_back(), None);
787            }
788        }
789    }
790
791    #[test]
792    fn test_divisors() {
793        let mut rng = Xorshift::default();
794        let pl = PrimeList::new(20000);
795        for n in (1..1000).chain(rng.random_iter(1..=20000000).take(100)) {
796            assert_eq!(pl.divisors(n), divisors(n));
797        }
798    }
799}
competitive/math/prime_list.rs

competitive/math/
prime_list.rs
