AoC_2020/src/day11.rs

use aoc_runner_derive::{aoc, aoc_generator};
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum Seat {
    Floor,
    Empty,
    Occupied,
}

pub struct Data(Vec<Vec<Seat>>);

impl Data {
    fn count_nb_p1(&self, x: usize, y: usize) -> usize {
        let mut ret = 0;
        for dy in &[-1i64, 0, 1] {
            for dx in &[-1i64, 0, 1] {
                if *dx == 0 && *dy == 0 {
                    continue;
                }
                let x = ((x as i64) + dx) as usize;
                let y = ((y as i64) + dy) as usize;
                if let Some(v) = self.0.get(y).map(|v| v.get(x)).flatten() {
                    if let Seat::Occupied = v {
                        ret += 1;
                    }
                }
            }
        }
        return ret;
    }

    fn count_nb_p2(&self, x: usize, y: usize) -> usize {
        let mut ret = 0;
        for dy in &[-1i64, 0, 1] {
            for dx in &[-1i64, 0, 1] {
                if *dx == 0 && *dy == 0 {
                    continue;
                }
                for l in 1.. {
                    let x = (x as i64) + dx * l;
                    let y = (y as i64) + dy * l;
                    if x < 0 || y < 0 {
                        break;
                    }
                    if let Some(v) = self.0.get(y as usize).map(|v| v.get(x as usize)).flatten() {
                        match v {
                            Seat::Occupied => {
                                ret += 1;
                                break;
                            }
                            Seat::Empty => {
                                break;
                            }
                            _ => (),
                        }
                    } else {
                        break;
                    }
                }
            }
        }
        ret
    }

    fn update_p1(&mut self) -> bool {
        let mut data = self.0.clone();
        for (y, row) in self.0.iter().enumerate() {
            for (x, seat) in row.iter().enumerate() {
                let occupied = self.count_nb_p1(x, y);
                match seat {
                    Seat::Floor => {}
                    Seat::Empty => {
                        if occupied == 0 {
                            data[y][x] = Seat::Occupied;
                        }
                    }
                    Seat::Occupied => {
                        if occupied > 3 {
                            data[y][x] = Seat::Empty;
                        }
                    }
                }
            }
        }
        let changed = data != self.0;
        self.0 = data;
        changed
    }

    fn update_p2(&mut self) -> bool {
        let mut data = self.0.clone();
        for (y, row) in self.0.iter().enumerate() {
            for (x, seat) in row.iter().enumerate() {
                let occupied = self.count_nb_p2(x, y);
                match seat {
                    Seat::Floor => {}
                    Seat::Empty => {
                        if occupied == 0 {
                            data[y][x] = Seat::Occupied;
                        }
                    }
                    Seat::Occupied => {
                        if occupied > 4 {
                            data[y][x] = Seat::Empty;
                        }
                    }
                }
            }
        }
        let changed = data != self.0;
        self.0 = data;
        changed
    }
}

#[aoc_generator(day11)]
pub fn input_generator(input: &str) -> Data {
    let ret: Vec<Vec<_>> = input
        .trim()
        .lines()
        .map(|line| {
            line.trim()
                .chars()
                .map(|c| match c {
                    'L' => Seat::Empty,
                    '#' => Seat::Occupied,
                    '.' => Seat::Floor,
                    c => panic!("invalid character: {:?}", c),
                })
                .collect()
        })
        .collect();
    for line in &ret {
        if ret[0].len() != line.len() {
            panic!("Non rectangular input!");
        }
    }
    Data(ret)
}

#[aoc(day11, part1)]
pub fn solve_part1(input: &Data) -> usize {
    let mut input = Data(input.0.clone());
    let mut n = 0;
    while input.update_p1() {}
    for row in &input.0 {
        for seat in row {
            if *seat == Seat::Occupied {
                n += 1;
            }
        }
    }
    return n;
}

#[aoc(day11, part2)]
pub fn solve_part2(input: &Data) -> usize {
    let mut input = Data(input.0.clone());
    while input.update_p2() {}

    let mut n = 0;
    for row in &input.0 {
        for seat in row {
            if *seat == Seat::Occupied {
                n += 1;
            }
        }
    }
    return n;
}