competitive/math/
quotient_array.rs

1use super::{Group, Invertible, One, Ring, Zero, with_prime_list};
2use std::ops::{Index, IndexMut};
3
4/// store with index $\{\lfloor\frac{n}{i}\rfloor \mid i=1,2,\ldots,n\}$
5#[derive(Debug, Clone)]
6pub struct QuotientArray<T> {
7    n: u64,
8    isqrtn: u64,
9    data: Vec<T>,
10}
11
12impl<T> QuotientArray<T>
13where
14    T: Zero,
15{
16    pub fn zeros(n: u64) -> Self {
17        Self::from_fn(n, |_| T::zero())
18    }
19}
20
21impl<T> QuotientArray<T> {
22    pub fn index_iter(n: u64, isqrtn: u64) -> impl Iterator<Item = u64> {
23        (1..=isqrtn)
24            .map(move |i| n / i)
25            .chain((1..n / isqrtn).rev())
26    }
27
28    pub fn map<U>(&self, f: impl FnMut(&T) -> U) -> QuotientArray<U> {
29        let data = self.data.iter().map(f).collect();
30        QuotientArray {
31            n: self.n,
32            isqrtn: self.isqrtn,
33            data,
34        }
35    }
36
37    pub fn quotient_index(&self, i: u64) -> usize {
38        assert!(
39            i <= self.n,
40            "index out of bounds: the len is {} but the index is {}",
41            self.n,
42            i
43        );
44        assert_ne!(i, 0, "index out of bounds: the index is 0");
45        if i <= self.isqrtn {
46            self.data.len() - i as usize
47        } else {
48            (self.n / i) as usize - 1
49        }
50    }
51
52    pub fn from_fn(n: u64, f: impl FnMut(u64) -> T) -> Self {
53        let isqrtn = (n as f64).sqrt().floor() as u64;
54        let data = Self::index_iter(n, isqrtn).map(f).collect();
55        Self { n, isqrtn, data }
56    }
57
58    /// convert $\sum_{i\leq n} f(i)$ to $\sum_{i\leq n, i\text{ is prime}} f(i)$
59    ///
60    /// constraints: $\mathrm{mul_p}(f(x))=f(px)$
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    }
116}
117
118impl<T> Index<u64> for QuotientArray<T> {
119    type Output = T;
120    fn index(&self, i: u64) -> &Self::Output {
121        unsafe { self.data.get_unchecked(self.quotient_index(i)) }
122    }
123}
124
125impl<T> IndexMut<u64> for QuotientArray<T> {
126    fn index_mut(&mut self, index: u64) -> &mut Self::Output {
127        let i = self.quotient_index(index);
128        unsafe { self.data.get_unchecked_mut(i) }
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135    use crate::{
136        algebra::{AddMulOperation, AdditiveOperation, ArrayOperation},
137        math::{PrimeList, PrimeTable},
138        tools::Xorshift,
139    };
140
141    #[test]
142    fn prime_count() {
143        let mut rng = Xorshift::default();
144        let pl = PrimeList::new(100_000);
145        for n in 1..=100 {
146            let n = if n <= 10 { n } else { rng.random(1..10_000) };
147            let qa = QuotientArray::from_fn(n, |i| i as i64 - 1)
148                .lucy_dp::<AdditiveOperation<_>>(|x, _p| x);
149            assert_eq!(pl.primes_lte(n as u32).count(), qa[n] as usize);
150        }
151    }
152
153    #[test]
154    fn divisor_sum() {
155        let mut rng = Xorshift::default();
156        let pt = PrimeTable::new(10_000);
157        for n in 1..=100 {
158            let n = if n <= 10 { n } else { rng.random(1..10_000) };
159            let qa = QuotientArray::from_fn(n, |i| [i as i64, i as i64 * (i as i64 + 1) / 2])
160                .map(|[x, y]| [x - 1, y - 1])
161                .lucy_dp::<ArrayOperation<AdditiveOperation<_>, 2>>(|[x, y], p| [x, y * p as i64])
162                .map(|[x, y]| x + y)
163                .min_25_sieve::<AddMulOperation<_>>(|p, c| {
164                    let mut x = 1;
165                    let mut s = 1;
166                    for _ in 0..c {
167                        x *= p as i64;
168                        s += x;
169                    }
170                    s
171                });
172            assert_eq!(
173                (1..=n)
174                    .flat_map(|i| pt.divisors(i as _))
175                    .map(|d| d as u64)
176                    .sum::<u64>(),
177                qa[n] as u64
178            );
179        }
180    }
181}
competitive/math/quotient_array.rs

competitive/math/
quotient_array.rs
