Rust

I hate when you need to look at the data to solve a simpler subclass of the problem.

use num::Integer;

fn one_round(columns: &mut [i64], forward: bool) -> bool {
    let mut modified = false;
    for i in 1..columns.len() {
        if forward {
            if columns[i - 1] > columns[i] {
                columns[i - 1] -= 1;
                columns[i] += 1;
                modified = true;
            }
        } else if columns[i - 1] < columns[i] {
            columns[i - 1] += 1;
            columns[i] -= 1;
            modified = true;
        }
    }
    modified
}

pub fn solve_part_1(input: &str) -> String {
    let mut columns = input
        .lines()
        .map(|l| l.parse().unwrap())
        .collect::<Vec<i64>>();
    let mut rounds_remaining = 10;
    while rounds_remaining > 0 {
        if one_round(&mut columns, true) {
            rounds_remaining -= 1;
        } else {
            break;
        }
    }
    while rounds_remaining > 0 {
        if one_round(&mut columns, false) {
            rounds_remaining -= 1;
        } else {
            break;
        }
    }
    columns
        .iter()
        .enumerate()
        .map(|(column_idx, count)| (column_idx + 1) as i64 * *count)
        .sum::<i64>()
        .to_string()
}

pub fn solve_part_2(input: &str) -> String {
    let mut columns = input
        .lines()
        .map(|l| l.parse().unwrap())
        .collect::<Vec<i64>>();
    let mut total_rounds = 0;
    loop {
        if one_round(&mut columns, true) {
            total_rounds += 1;
        } else {
            break;
        }
    }
    loop {
        if one_round(&mut columns, false) {
            total_rounds += 1;
        } else {
            break;
        }
    }
    total_rounds.to_string()
}

pub fn solve_part_3(input: &str) -> String {
    let mut columns = input
        .lines()
        .map(|l| l.parse().unwrap())
        .collect::<Vec<i64>>();
    let invariant_sum = columns.iter().sum::<i64>();
    let mut total_rounds = 0i64;
    loop {
        let first_gap = (0..columns.len() - 1).find(|&i| columns[i].abs_diff(columns[i + 1]) > 1);
        let last_gap = (0..columns.len() - 1)
            .filter(|&i| columns[i].abs_diff(columns[i + 1]) > 1)
            .next_back();
        if let (Some(first_gap), Some(last_gap)) = (first_gap, last_gap)
            && (last_gap > first_gap)
        {
            let amount_to_fill: i64 = columns
                .iter()
                .take(first_gap + 1)
                .map(|&x| columns[first_gap + 1] - x)
                .sum();
            let amount_available: i64 = columns
                .iter()
                .skip(last_gap + 1)
                .map(|&x| x - columns[last_gap])
                .sum();
            let amount_to_transfer = amount_to_fill.min(amount_available);
            let new_amount_left: i64 =
                columns.iter().take(first_gap + 1).sum::<i64>() + amount_to_transfer;
            let (left_base, remainder) = new_amount_left.div_rem(&((first_gap + 1) as i64));
            for (i, new_value) in columns.iter_mut().enumerate().take(first_gap + 1) {
                *new_value = left_base
                    + if first_gap - i < remainder as usize {
                        1
                    } else {
                        0
                    };
            }
            let new_amount_right: i64 =
                columns.iter().skip(last_gap + 1).sum::<i64>() - amount_to_transfer;
            let (right_base, remainder) =
                new_amount_right.div_rem(&((columns.len() - last_gap - 1) as i64));
            for i in last_gap + 1..columns.len() {
                columns[i] = right_base
                    + if columns.len() - i <= remainder as usize {
                        1
                    } else {
                        0
                    };
            }
            assert_eq!(invariant_sum, columns.iter().sum::<i64>());
            total_rounds += amount_to_transfer;
        } else {
            break;
        }
    }
    loop {
        if one_round(&mut columns, false) {
            total_rounds += 1;
        } else {
            break;
        }
    }
    total_rounds.to_string()
}