haar_lib/math/
sparse_polynomial.rs1use crate::num::const_modint::*;
3
4#[derive(Clone, Debug, Default)]
6pub struct SparsePolynomial<const P: u32> {
7 pub(crate) data: Vec<(usize, ConstModInt<P>)>,
8}
9
10impl<const P: u32> SparsePolynomial<P> {
11 pub fn zero() -> Self {
13 Self { data: vec![] }
14 }
15
16 pub fn constant(a: ConstModInt<P>) -> Self {
18 if a.value() == 0 {
19 Self::zero()
20 } else {
21 Self { data: vec![(0, a)] }
22 }
23 }
24
25 pub fn add(&mut self, i: usize, x: ConstModInt<P>) {
26 for (j, y) in self.data.iter_mut() {
27 if i == *j {
28 *y += x;
29 return;
30 }
31 }
32 self.data.push((i, x));
33 }
34
35 pub fn from_vec(mut a: Vec<(usize, ConstModInt<P>)>) -> Self {
36 a.sort_by_key(|a| a.0);
37 Self { data: a }
38 }
39
40 pub fn coeff_of(&self, i: usize) -> ConstModInt<P> {
42 self.data
43 .iter()
44 .find(|(j, _)| *j == i)
45 .map_or(0.into(), |(_, x)| *x)
46 }
47
48 pub fn differential(&mut self) {
50 let a = self
51 .data
52 .iter()
53 .filter_map(|&(i, x)| {
54 if i == 0 {
55 None
56 } else {
57 Some((i - 1, x * i.into()))
58 }
59 })
60 .collect();
61
62 self.data = a;
63 }
64
65 pub fn integral(&mut self) {
67 let a = self
68 .data
69 .iter()
70 .map(|&(i, x)| (i + 1, x * ConstModInt::new(i as u32 + 1).inv()))
71 .collect();
72
73 self.data = a;
74 }
75
76 pub fn scale(&mut self, k: ConstModInt<P>) {
78 self.data.iter_mut().for_each(|(_, x)| *x *= k);
79 }
80
81 pub fn shift_lower(&mut self, k: usize) {
83 self.data = self
84 .data
85 .iter()
86 .filter_map(|&(i, x)| if i < k { None } else { Some((i - k, x)) })
87 .collect();
88 }
89
90 pub fn shift_higher(&mut self, k: usize) {
92 self.data.iter_mut().for_each(|(i, _)| *i += k);
93 }
94}