haar_lib/math/polynomial/
sparse.rs

1//! 疎な多項式
2use std::collections::BTreeMap;
3
4use crate::math::prime_mod::PrimeMod;
5use crate::num::const_modint::*;
6
7/// 疎な多項式
8#[derive(Clone, Debug, Default)]
9pub struct SparsePolynomial<P: PrimeMod> {
10    data: BTreeMap<usize, ConstModInt<P>>,
11}
12
13impl<P: PrimeMod> SparsePolynomial<P> {
14    /// 零多項式を得る。
15    pub fn zero() -> Self {
16        Self {
17            data: BTreeMap::new(),
18        }
19    }
20
21    /// 定数項のみをもつ多項式を生成する。
22    pub fn constant(a: ConstModInt<P>) -> Self {
23        if a.value() == 0 {
24            Self::zero()
25        } else {
26            Self {
27                data: BTreeMap::from([(0, a)]),
28            }
29        }
30    }
31
32    /// $k x^i$となる項を足す。
33    pub fn add_one_term<T>(&mut self, i: usize, k: T)
34    where
35        ConstModInt<P>: From<T>,
36    {
37        *self.data.entry(i).or_default() += k.into();
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 = self.data.iter().map(|(&i, &x)| (i + k, x)).collect();
93    }
94
95    /// 項の次数と係数のペアへのイテレータを返す。
96    pub fn iter(&self) -> impl Iterator<Item = (&usize, &ConstModInt<P>)> {
97        self.data.iter()
98    }
99}
100
101impl<P: PrimeMod, T> From<Vec<(usize, T)>> for SparsePolynomial<P>
102where
103    ConstModInt<P>: From<T>,
104{
105    fn from(value: Vec<(usize, T)>) -> Self {
106        let mut ret = Self::zero();
107        for (i, x) in value {
108            ret.add_one_term(i, x);
109        }
110        ret
111    }
112}
haar_lib/math/polynomial/sparse.rs

haar_lib/math/polynomial/
sparse.rs
