struct LeastSquares {
    m0: f64,
    m1: f64,
    m2: f64,
    b0: f64,
    b1: f64
}

impl LeastSquares {
    fn new() -> LeastSquares {
        LeastSquares { m0: 0.0, m1: 0.0, m2: 0.0, b0: 0.0, b1: 0.0 }
    }

    fn insert(&mut self, x: f64, y: f64) {
        self.m0 += 1.0;
        self.m1 += x;
        self.m2 += x * x;
        self.b0 += y;
        self.b1 += x * y;
    }

    fn solve(self) -> (f64, f64) {
        // a m2 + b m1 = b1
        // a m1 + b m0 = b0
        let u = self.m0 * self.b1 - self.m1 * self.b0;
        let v = self.m2 * self.b0 - self.m1 * self.b1;
        let w = self.m2 * self.m0 - self.m1 * self.m1;
        return (u / w, v / w);
    }
}

fn cluster(mut a: f64, mut b: f64, xs: &[f64]) -> Vec<i64> {
    let mut is = vec![0i64; xs.len()];
    let mut changed = true;
    while changed {
        changed = false;
        let mut ls = LeastSquares::new();
        for (x, cur_i) in xs.iter().zip(is.iter_mut()) {
            let new_i = ((x - b) / a).round() as i64;
            ls.insert(new_i as f64, *x);
            changed = changed || new_i != *cur_i;
            *cur_i = new_i;
        }

        let (new_a, new_b) = ls.solve();
        a = new_a; b = new_b;
    }

    return is;
}

fn main() {
    let is = cluster(3.0, 1.5, &[1.5, 4.6, 7.1, 10.9]);
    println!("{:?}", is);
}