fn almost_magic_square(a: &[i64]) -> Vec<i64> {
    let mut grid: Vec<Vec<i64>> = a.chunks(3).map(|c| c.to_vec()).collect();
    let mut row_sum = [0i64; 3];
    let mut col_sum = [0i64; 3];

    for i in 0..3 {
        for j in 0..3 {
            row_sum[i] += grid[i][j];
            col_sum[i] += grid[j][i];
        }
    }

    let mut max_sum = 0;
    for k in 0..3 {
        max_sum = max_sum.max(row_sum[k]);
        max_sum = max_sum.max(col_sum[k]);
    }

    let (mut i, mut j) = (0usize, 0usize);
    while i < 3 && j < 3 {
        let diff = (max_sum - row_sum[i]).min(max_sum - col_sum[j]);
        grid[i][j] += diff;
        row_sum[i] += diff;
        col_sum[j] += diff;

        if row_sum[i] == max_sum {
            i += 1;
        }
        if col_sum[j] == max_sum {
            j += 1;
        }
    }

    grid.into_iter().flatten().collect()
}