AoC2021/src/day07.rs

use std::collections::HashMap;
use crate::Day;
use anyhow::Result;

pub struct Day07(HashMap<i32, i32>);

impl Day for Day07 {
    fn init(content: String) -> Result<Self> {
        let mut hm = HashMap::new();
        for n in &content.split(',').map(str::parse).collect::<Result<Vec<_>, std::num::ParseIntError>>()? {
            let entry = hm.entry(*n).or_insert(0);
            *entry += 1;
        }
        Ok(Self(hm))
    }

    fn part1(&self) -> Result<String> {
        let m = self.0.keys().copied().map(
            |x| self.0.iter().map(
                |(n, count)| (x - *n).abs() * (*count)
            ).sum::<i32>()
        ).min().ok_or_else(|| anyhow::Error::msg("Could not find min"))?;

        Ok(format!("{}", m))
    }

    fn part2(&self) -> Result<String> {
        let max = *self.0.keys().max().ok_or_else(|| anyhow::Error::msg("Could not find max"))?;
        let m = (0..=max).map(
            |x| self.0.iter().map(
                |(n, count)| {
                    let n = (x - *n).abs();
                    (n * (n + 1) / 2) * (*count)
                }
            ).sum::<i32>()
        ).min().ok_or_else(|| anyhow::Error::msg("Could not find min"))?;
        Ok(format!("{}", m))
    }
}

#[cfg(test)]
mod tests {
    use crate::day07::Day07;
    use crate::day::Day;
    use anyhow::Result;
    use crate::day_tests;

    const INPUT: &str = "16,1,2,0,4,2,7,1,2,14";

    #[test]
    fn part1_test() -> Result<()> {
        let d = Day07::init(INPUT.to_string())?;
        assert_eq!("37", d.part1()?);
        Ok(())
    }

    #[test]
    fn part2_test() -> Result<()> {
        let d = Day07::init(INPUT.to_string())?;
        assert_eq!("168", d.part2()?);
        Ok(())
    }

    day_tests!(Day07, "07", "357353", "104822130");
}