const fn table_k<const P: u32>() -> usizeExamples found in repository?
crates/competitive/src/math/fast_prime_mod.rs (line 24)
23const fn table_len<const P: u32>() -> usize {
24 2 * table_k::<P>() + 1
25}
26
27/// Fast online inverse and power queries for a prime modulus.
28pub struct FastPrimeMod<const P: u32, const BUILD_INV: bool = true, const BUILD_POW: bool = true> {
29 root: u32,
30 pow_lo: Box<[u32]>,
31 pow_hi: Box<[u32]>,
32 frac: Box<[u32]>,
33 log: Box<[u32]>,
34 inv: Box<[u32]>,
35}
36
37impl<const P: u32, const BUILD_INV: bool, const BUILD_POW: bool>
38 FastPrimeMod<P, BUILD_INV, BUILD_POW>
39{
40 /// Builds the tables enabled by the const generic mode.
41 ///
42 /// # Panics
43 ///
44 /// Panics if `P` is not an odd prime or if `P >= 2^30`.
45 pub fn new() -> Self {
46 assert!(
47 BUILD_INV || BUILD_POW,
48 "at least one of BUILD_INV or BUILD_POW must be true"
49 );
50 assert!(P < 1 << 30, "P must be less than 2^30");
51 assert!(
52 P % 2 == 1 && miller_rabin(P as u64),
53 "P must be an odd prime"
54 );
55
56 let (root, pow_lo, pow_hi, log) = if BUILD_POW {
57 let root = if P == 998_244_353 {
58 3
59 } else {
60 primitive_root(P as u64) as u32
61 };
62 let (pow_lo, pow_hi) = build_pow::<P>(root);
63 let log = build_log::<P>(root, &pow_lo, &pow_hi);
64 (root, pow_lo, pow_hi, log)
65 } else {
66 (
67 0,
68 Vec::new().into_boxed_slice(),
69 Vec::new().into_boxed_slice(),
70 Vec::new().into_boxed_slice(),
71 )
72 };
73 let inv = if BUILD_INV {
74 build_inv::<P>()
75 } else {
76 Vec::new().into_boxed_slice()
77 };
78 let frac = if is_direct_table_mod::<P>() {
79 Vec::new().into_boxed_slice()
80 } else {
81 build_frac::<P>()
82 };
83 Self {
84 root,
85 pow_lo,
86 pow_hi,
87 frac,
88 log,
89 inv,
90 }
91 }
92
93 /// Returns the prime modulus.
94 #[inline]
95 pub fn modulus(&self) -> u32 {
96 P
97 }
98
99 #[inline(always)]
100 fn small_fraction(&self, x: u32) -> (usize, u32) {
101 let k = table_k::<P>();
102 if is_direct_table_mod::<P>() {
103 debug_assert!(1 <= x && x < P);
104 return (k + x as usize, 1);
105 }
106 let packed = self.frac[(x >> FRAC_SHIFT) as usize];
107 let a = packed >> 16;
108 let b = packed & 0xffff;
109 let t = x.wrapping_mul(b).wrapping_sub(a.wrapping_mul(P));
110 debug_assert!({
111 let t = x as i64 * b as i64 - a as i64 * P as i64;
112 t != 0 && -(k as i64) <= t && t <= k as i64
113 });
114 ((k as u32).wrapping_add(t) as usize, b)
115 }
116}
117
118impl<const P: u32, const BUILD_POW: bool> FastPrimeMod<P, true, BUILD_POW> {
119 /// Returns `x^{-1} mod P`.
120 ///
121 /// # Panics
122 ///
123 /// Panics unless `1 <= x < P`.
124 #[inline]
125 pub fn inverse(&self, x: u32) -> u32 {
126 assert!(1 <= x && x < P);
127 let (i, b) = self.small_fraction(x);
128 mul_mod_raw::<P>(self.inv[i], b)
129 }
130}
131
132impl<const P: u32, const BUILD_INV: bool> FastPrimeMod<P, BUILD_INV, true> {
133 /// Returns the primitive root used by this table.
134 #[inline]
135 pub fn primitive_root(&self) -> u32 {
136 self.root
137 }
138
139 /// Returns `a^exp mod P`.
140 ///
141 /// `0^0` is defined as `1`.
142 ///
143 /// # Panics
144 ///
145 /// Panics unless `a < P`.
146 #[inline]
147 pub fn pow(&self, a: u32, exp: u64) -> u32 {
148 assert!(a < P);
149 if a == 0 {
150 return if exp == 0 { 1 } else { 0 };
151 }
152 let ord = (P - 1) as u64;
153 self.pow_nonzero_reduced(a, (exp % ord) as u32)
154 }
155
156 /// Returns `a^exp_mod mod P` for a non-zero base and a reduced exponent.
157 ///
158 /// # Panics
159 ///
160 /// Panics unless `1 <= a < P` and `exp_mod < P - 1`.
161 #[inline]
162 pub fn pow_nonzero_reduced(&self, a: u32, exp_mod: u32) -> u32 {
163 assert!(1 <= a && a < P);
164 assert!(exp_mod < P - 1);
165 let exp = (self.log_r(a) as u64 * exp_mod as u64 % (P - 1) as u64) as u32;
166 self.pow_root_reduced(exp)
167 }
168
169 /// Returns `r^exp_mod mod P`, where `r` is this table's primitive root.
170 ///
171 /// # Panics
172 ///
173 /// Panics unless `exp_mod < P - 1`.
174 #[inline]
175 pub fn pow_root_reduced(&self, exp_mod: u32) -> u32 {
176 assert!(exp_mod < P - 1);
177 pow_root_raw::<P>(exp_mod, &self.pow_lo, &self.pow_hi)
178 }
179
180 #[inline]
181 fn log_r(&self, x: u32) -> u32 {
182 let (i, b) = self.small_fraction(x);
183 let k = table_k::<P>();
184 let ord = P - 1;
185 self.log[i] + ord - self.log[k + b as usize]
186 }
187}
188
189impl<const P: u32, const BUILD_INV: bool, const BUILD_POW: bool> Default
190 for FastPrimeMod<P, BUILD_INV, BUILD_POW>
191{
192 fn default() -> Self {
193 Self::new()
194 }
195}
196
197fn build_pow<const P: u32>(root: u32) -> (Box<[u32]>, Box<[u32]>) {
198 let mut pow_lo = vec![0; POW_BLOCK + 1].into_boxed_slice();
199 let mut pow_hi = vec![0; POW_BLOCK + 1].into_boxed_slice();
200 pow_lo[0] = 1;
201 pow_hi[0] = 1;
202 for i in 0..POW_BLOCK {
203 pow_lo[i + 1] = mul_mod_raw::<P>(pow_lo[i], root);
204 }
205 let block_power = pow_lo[POW_BLOCK];
206 for i in 0..POW_BLOCK {
207 pow_hi[i + 1] = mul_mod_raw::<P>(pow_hi[i], block_power);
208 }
209 (pow_lo, pow_hi)
210}
211
212fn build_inv<const P: u32>() -> Box<[u32]> {
213 let k = table_k::<P>();
214 let mut inv = vec![0; table_len::<P>()].into_boxed_slice();
215 inv[k + 1] = 1;
216 for i in 2..=k {
217 let q = P.div_ceil(i as u32);
218 let r = i as u32 * q - P;
219 inv[k + i] = mul_mod_raw::<P>(inv[k + r as usize], q);
220 }
221 for i in 1..=k {
222 inv[k - i] = P - inv[k + i];
223 }
224 inv
225}
226
227fn build_log<const P: u32>(root: u32, pow_lo: &[u32], pow_hi: &[u32]) -> Box<[u32]> {
228 let k = table_k::<P>();
229 let ord = P - 1;
230 let mut lpf = vec![0; k + 1].into_boxed_slice();
231 let mut primes = vec![];
232 lpf[1] = 1;
233 for i in 2..=k {
234 if lpf[i] == 0 {
235 lpf[i] = i as u32;
236 primes.push(i as u32);
237 }
238 for &p in primes.iter() {
239 let p = p as usize;
240 if p > lpf[i] as usize || p > k / i {
241 break;
242 }
243 lpf[i * p] = p as u32;
244 }
245 }
246
247 let baby_size = (BSGS_SIZE as u32).min(ord);
248 let mut baby = U32Map::new(baby_size as usize);
249 let mut pw = 1;
250 for i in 0..baby_size {
251 baby.insert(pw, i);
252 pw = mul_mod_raw::<P>(pw, root);
253 }
254 let q = pow_root_raw::<P>(ord - baby_size, pow_lo, pow_hi);
255
256 let mut log = vec![0; table_len::<P>()].into_boxed_slice();
257 log[k + 1] = 0;
258 let mut rng = Xorshift::default();
259 let small_primes = [2, 3, 5, 7, 11, 13, 17, 19];
260 for i in 2..=k {
261 let p = lpf[i] as usize;
262 if p < i {
263 log[k + i] = add_mod(log[k + p], log[k + i / p], ord);
264 } else if i < 100 {
265 let mut x = i as u32;
266 let mut ans = 0;
267 loop {
268 if let Some(v) = baby.get(x) {
269 log[k + i] = ans + v;
270 break;
271 }
272 ans += baby_size;
273 x = mul_mod_raw::<P>(x, q);
274 }
275 } else if i > P as usize / i {
276 let j = (P as usize) / i;
277 let r = (P as usize) % i;
278 let x = add_mod(log[k + r], ord / 2, ord);
279 let y = log[k + j];
280 log[k + i] = if x >= y { x - y } else { x + ord - y };
281 } else {
282 loop {
283 let exp = rng.rand(ord as u64) as u32;
284 let mut ans = if exp == 0 { 0 } else { ord - exp };
285 let mut x = mul_mod_raw::<P>(i as u32, pow_root_raw::<P>(exp, pow_lo, pow_hi));
286 for q in small_primes {
287 while x.is_multiple_of(q) {
288 x /= q;
289 ans = add_mod(ans, log[k + q as usize], ord);
290 }
291 }
292 if x as usize >= k {
293 continue;
294 }
295 while (i as u32) < x && lpf[x as usize] < i as u32 {
296 let q = lpf[x as usize];
297 x /= q;
298 ans = add_mod(ans, log[k + q as usize], ord);
299 }
300 if 1 < x && x < i as u32 {
301 ans = add_mod(ans, log[k + x as usize], ord);
302 x = 1;
303 }
304 if x == 1 {
305 log[k + i] = ans;
306 break;
307 }
308 }
309 }
310 }
311 for i in 1..=k {
312 log[k - i] = add_mod(log[k + i], ord / 2, ord);
313 }
314 log
315}