haar_lib/math/
sparse_polynomial.rs

1//! 疎な多項式
2use crate::num::const_modint::*;
3
4/// 疎な多項式
5#[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    /// 零多項式を得る。
12    pub fn zero() -> Self {
13        Self { data: vec![] }
14    }
15
16    /// 定数項のみをもつ多項式を生成する。
17    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    /// $x^i$の係数を得る。
41    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    /// 多項式を微分する。
49    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    /// 多項式を積分する。
66    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    /// 多項式を`k`倍する。
77    pub fn scale(&mut self, k: ConstModInt<P>) {
78        self.data.iter_mut().for_each(|(_, x)| *x *= k);
79    }
80
81    /// 係数を`k`次だけ低次側にずらす。ただし、負の次数の項は無視する。
82    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    /// 係数を`k`次だけ高次側にずらす。
91    pub fn shift_higher(&mut self, k: usize) {
92        self.data.iter_mut().for_each(|(i, _)| *i += k);
93    }
94}