fit.rs

1const P: u128 = (1 << 61) - 1;23fn power(mut a: u128, mut e: u128) -> u128 {4    let mut r = 1;5    a %= P;6    while e > 0 {7        if e & 1 == 1 {8            r = r * a % P;9        }10        a = a * a % P;11        e >>= 1;12    }13    r14}1516fn massey(s: &[u128]) -> Vec<u128> {17    let (mut c, mut b) = (vec![1u128], vec![1u128]);18    let (mut l, mut m, mut last) = (0usize, 1usize, 1u128);19    for n in 0..s.len() {20        let mut d = s[n] % P;21        for i in 1..=l.min(c.len() - 1) {22            d = (d + c[i] * (s[n - i] % P)) % P;23        }24        if d == 0 {25            m += 1;26            continue;27        }28        let t = c.clone();29        let coef = d * power(last, P - 2) % P;30        if c.len() < b.len() + m {31            c.resize(b.len() + m, 0);32        }33        for i in 0..b.len() {34            c[i + m] = (c[i + m] + P - coef * b[i] % P) % P;35        }36        if 2 * l <= n {37            l = n + 1 - l;38            b = t;39            last = d;40            m = 1;41        } else {42            m += 1;43        }44    }45    c.resize(l + 1, 0);46    c47}4849pub struct Fit {50    pub poly: Vec<i128>,51    pub margin: i64,52}5354pub fn fit(s: &[u128]) -> Option<Fit> {55    let c = massey(s);56    let poly: Vec<i128> = c57        .iter()58        .map(|&x| {59            if x > P / 2 {60                x as i128 - P as i12861            } else {62                x as i12863            }64        })65        .collect();66    let l = poly.len() - 1;67    for n in l..s.len() {68        let mut acc = 0i128;69        for (i, &k) in poly.iter().enumerate() {70            acc = acc.checked_add(k.checked_mul(s[n - i] as i128)?)?;71        }72        if acc != 0 {73            return None;74        }75    }76    Some(Fit {77        poly,78        margin: s.len() as i64 - 2 * l as i64,79    })80}8182fn eval(poly: &[i128], r: i128) -> Option<i128> {83    poly.iter()84        .try_fold(0i128, |acc, &k| acc.checked_mul(r)?.checked_add(k))85}8687pub fn roots(poly: &[i128]) -> (Vec<i128>, Vec<i128>) {88    let mut rest = poly.to_vec();89    let mut found = Vec::new();90    for r in (0..=64i128).flat_map(|r| if r == 0 { vec![0] } else { vec![r, -r] }) {91        while rest.len() > 1 && eval(&rest, r) == Some(0) {92            let mut q = Vec::with_capacity(rest.len() - 1);93            let mut acc = 0i128;94            for &k in &rest[..rest.len() - 1] {95                acc = acc * r + k;96                q.push(acc);97            }98            rest = q;99            found.push(r);100        }101    }102    found.sort_by(|a, b| b.cmp(a));103    (found, rest)104}105106pub fn show(poly: &[i128]) -> String {107    let d = poly.len() - 1;108    let mut out = String::new();109    for (i, &k) in poly.iter().enumerate() {110        if k == 0 {111            continue;112        }113        let e = d - i;114        let sign = if k < 0 {115            " - "116        } else if out.is_empty() {117            ""118        } else {119            " + "120        };121        let a = k.abs();122        let lead = if out.is_empty() && k < 0 {123            "-".to_string()124        } else {125            sign.to_string()126        };127        let body = match (a, e) {128            (_, 0) => format!("{a}"),129            (1, 1) => "x".to_string(),130            (1, _) => format!("x^{e}"),131            (_, 1) => format!("{a} x"),132            _ => format!("{a} x^{e}"),133        };134        out.push_str(&lead);135        out.push_str(&body);136    }137    out138}