AoC2020/src/day11.rs

use std::io::BufRead;

#[derive(Clone, Debug, PartialEq)]
enum Cell {
    Empty,
    Taken,
    Floor,
}

struct Config {
    max_steps: Option<usize>,
    max_taken: usize,
}

type FloorMap = Vec<Vec<Cell>>;
type UpdateEvent = (usize, usize, Cell);

const CONFIG_PART1: Config = Config {
    max_steps: Some(1),
    max_taken: 4,
};

const CONFIG_PART2: Config = Config {
    max_steps: None,
    max_taken: 5,
};

fn parse<F: BufRead> (input: F) -> FloorMap {
    input.lines()
        .map(|line| line.unwrap().chars().map(|c| match c {
            'L' => Cell::Empty,
            '#' => Cell::Taken,
            '.' => Cell::Floor,
            _ => unreachable!()
        }).collect::<Vec<_>>())
        .collect()
}

fn get_update(floormap: &FloorMap, line_idx: usize, col_idx: usize, config: &Config) -> Option<UpdateEvent> {
    let mut count_taken = 0;
    let current_cell = &floormap[line_idx][col_idx];

    for l in -1..2 {
        for c in -1..2 {
            if c == 0 && l == 0 {
                continue
            }
            let mut steps = 1;
            loop {
                if let Some(max_steps) = config.max_steps {
                    if steps > max_steps {
                        break;
                    }
                }

                let other_line_idx = line_idx as isize + l * steps as isize;
                let other_col_idx = col_idx as isize + c * steps as isize;

                if other_col_idx < 0 || other_line_idx < 0 {
                    break;
                }

                let cell = floormap
                    .get(other_line_idx as usize)
                    .map(|line| line.get(other_col_idx as usize))
                    .flatten();

                match cell {
                    Some(Cell::Taken) => {
                        count_taken += 1;
                        break
                    },
                    Some(Cell::Empty) | None => break,
                    _ => (),
                }

                steps += 1;
            }
        }
    }
    if *current_cell == Cell::Empty && count_taken == 0 {
        Some((line_idx, col_idx, Cell::Taken))
    } else if *current_cell == Cell::Taken && count_taken >= config.max_taken {
        Some((line_idx, col_idx, Cell::Empty))
    } else {
        None
    }
}

fn get_changes_floor_map(floormap: &FloorMap, config: &Config) -> Vec<UpdateEvent> {
    floormap.iter()
        .enumerate()
        .flat_map(|(line_idx, line)| {
            line.iter()
                .enumerate()
                .filter_map(move |(col_idx, cell)| match cell {
                    Cell::Floor => None,
                    _ => get_update(floormap, line_idx, col_idx, config)
                })
        }).collect()
}

fn update_floor_map(floormap: &mut FloorMap, mut updates: Vec<UpdateEvent>) {
    for (line_idx, col_idx, cell) in updates.drain(..) {
        floormap[line_idx][col_idx] = cell;
    }
}

fn tick(floormap: &mut FloorMap, config: &Config) -> bool {
    let updates = get_changes_floor_map(floormap, config);
    if updates.len() > 0 {
        update_floor_map(floormap, updates);
        true
    } else {
        false
    }
}

fn simulate_floormap(mut floormap: FloorMap, config: &Config) -> usize {
    while tick(&mut floormap, config) {}
    floormap.iter().map(|line| line.iter().filter(|&cell| *cell == Cell::Taken).count()).sum()
}

pub fn part1<F: BufRead> (input: F) {
    println!("{}", simulate_floormap(parse(input), &CONFIG_PART1));
}

pub fn part2<F: BufRead> (input: F) {
    println!("{}", simulate_floormap(parse(input), &CONFIG_PART2));
}

#[cfg(test)]
mod test {
    use super::*;
    const INPUT: &str = r#"L.LL.LL.LL
LLLLLLL.LL
L.L.L..L..
LLLL.LL.LL
L.LL.LL.LL
L.LLLLL.LL
..L.L.....
LLLLLLLLLL
L.LLLLLL.L
L.LLLLL.LL"#;

    #[test]
    pub fn test_tick() {
        let expected = r#"#.LL.L#.##
#LLLLLL.L#
L.L.L..L..
#LLL.LL.L#
#.LL.LL.LL
#.LLLL#.##
..L.L.....
#LLLLLLLL#
#.LLLLLL.L
#.#LLLL.##"#.as_bytes();

        let mut floormap = parse(INPUT.as_bytes());
        tick(&mut floormap, &CONFIG_PART1);
        tick(&mut floormap, &CONFIG_PART1);
        assert_eq!(parse(expected), floormap);
    }

    #[test]
    pub fn test_tick2() {
        let expected = r#"#.LL.LL.L#
#LLLLLL.LL
L.L.L..L..
LLLL.LL.LL
L.LL.LL.LL
L.LLLLL.LL
..L.L.....
LLLLLLLLL#
#.LLLLLL.L
#.LLLLL.L#"#.as_bytes();

        let mut floormap = parse(INPUT.as_bytes());
        tick(&mut floormap, &CONFIG_PART1);
        tick(&mut floormap, &CONFIG_PART2);
        assert_eq!(parse(expected), floormap);
    }

    #[test]
    pub fn test_simulate1() {
        assert_eq!(simulate_floormap(parse(INPUT.as_bytes()), &CONFIG_PART1), 37);
    }

    #[test]
    pub fn test_simulate2() {
        assert_eq!(simulate_floormap(parse(INPUT.as_bytes()), &CONFIG_PART2), 26);
    }
}