competitive_programming_rs/math/
lagrange_interpolation.rs

1pub fn lagrange_interpolation<T>(xs: &[T], ys: &[T], one: T, zero: T) -> Vec<T>
2where
3    T: Clone
4        + Copy
5        + std::ops::Sub<Output = T>
6        + std::ops::Mul<Output = T>
7        + std::ops::Div<Output = T>
8        + std::ops::AddAssign
9        + std::ops::SubAssign
10        + std::ops::MulAssign,
11{
12    let n = xs.len();
13
14    let mut all_c = vec![zero; n + 1];
15    all_c[0] = one;
16    for &x in xs.iter() {
17        let mut next = vec![zero; n + 1];
18        next[1..(n + 1)].clone_from_slice(&all_c[..n]);
19        for j in 0..n {
20            next[j] -= x * all_c[j];
21        }
22        all_c = next;
23    }
24
25    let mut c = vec![zero; n];
26    for i in 0..n {
27        let mut qi = one;
28        for j in 0..n {
29            if i == j {
30                continue;
31            }
32            qi *= xs[i] - xs[j];
33        }
34
35        let ri = ys[i] / qi;
36        let mut tmp_c = all_c.clone();
37        for j in (0..n).rev() {
38            c[j] += ri * tmp_c[j + 1];
39            let next_c = tmp_c[j + 1] * xs[i];
40            tmp_c[j] += next_c;
41        }
42    }
43    c
44}
45
46#[cfg(test)]
47mod tests {
48    use super::*;
49    use crate::math::mod_int::mod_int::{set_mod_int, ModInt};
50    use rand::{thread_rng, Rng};
51
52    const MOD: i64 = 1_000_000_007;
53
54    #[test]
55    fn test_lagrange_interpolation() {
56        set_mod_int(MOD);
57        let mut rng = thread_rng();
58
59        let n = 500;
60        for _ in 0..10 {
61            let mut xs = vec![];
62            let mut ys = vec![];
63            for _ in 0..n {
64                xs.push(ModInt::from(rng.gen_range(0, MOD)));
65                ys.push(ModInt::from(rng.gen_range(0, MOD)));
66            }
67
68            let c = lagrange_interpolation(&xs, &ys, ModInt::from(1), ModInt::from(0));
69
70            for i in 0..n {
71                let mut y = ModInt::from(0);
72                let x = xs[i];
73                for i in 0..n {
74                    y += x.pow(i as i64) * c[i];
75                }
76                assert_eq!(y.value(), ys[i].value());
77            }
78        }
79    }
80}
competitive_programming_rs/math/lagrange_interpolation.rs

competitive_programming_rs/math/
lagrange_interpolation.rs
