adventofcode-2020/day-10/src/main.rs

use std::convert::TryInto;
use std::fs;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("input")?;
    let jolts = parse_jolts(&input)?;

    // Part 1
    let result = chain_all(&jolts);
    println!("{}", result.0 * result.1);

    // Part 2
    let result = count_possible_chains(&jolts)?;
    println!("{}", result);

    Ok(())
}

fn count_possible_chains(jolts: &Vec<u64>) -> Result<u64, std::num::TryFromIntError> {
    let mut jolts = jolts.to_vec();
    jolts.push(0);
    jolts.sort();
    jolts.push(jolts.last().expect("Need some jolts") + 3);

    let window_lengths: Vec<usize> = jolts
        .windows(2)
        .scan(0, |length, window| {
            if let [previous, current, _empty @ ..] = window {
                if *current - *previous < 3 {
                    *length += 1;
                    Some(None)
                } else {
                    let full_length = *length;
                    *length = 0;
                    Some(Some(full_length))
                }
            } else {
                unreachable!("Window size is 2");
            }
        })
        .filter_map(|option| option)
        .collect();

    fn permutations(length: u32) -> u64 {
        if length == 0 {
            1
        } else {
            2u64.pow(length - 1) - if length > 3 { 2u64.pow(length - 4) } else { 0 }
        }
    }

    window_lengths
        .iter()
        .map(|length| Ok(permutations((*length).try_into()?)))
        .product()
}

fn chain_all(jolts: &Vec<u64>) -> (usize, usize) {
    let mut jolts = jolts.to_vec();
    jolts.push(0);
    jolts.sort();
    jolts.push(jolts.last().expect("Need some jolts") + 3);

    let jolt_jumps: Vec<u64> = jolts
        .windows(2)
        .map(|window| {
            if let [first, second, _empty @ ..] = window {
                second - first
            } else {
                unreachable!("Window size is 2");
            }
        })
        .collect();

    let jump_1s = jolt_jumps.iter().filter(|e| **e == 1).count();
    let jump_3s = jolt_jumps.iter().filter(|e| **e == 3).count();

    (jump_1s, jump_3s)
}

fn parse_jolts(input: &str) -> Result<Vec<u64>, std::num::ParseIntError> {
    input.lines().map(|line| line.parse()).collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use indoc::indoc;

    #[test]
    fn test_simple() -> Result<(), Box<dyn std::error::Error>> {
        let input = indoc!(
            "16
             10
             15
             5
             1
             11
             7
             19
             6
             12
             4
            "
        );

        let jolts = parse_jolts(input)?;
        let result = chain_all(&jolts);

        assert_eq!(result.0 * result.1, 7 * 5);
        Ok(())
    }

    #[test]
    fn test_simple2() -> Result<(), Box<dyn std::error::Error>> {
        let input = indoc!(
            "28
             33
             18
             42
             31
             14
             46
             20
             48
             47
             24
             23
             49
             45
             19
             38
             39
             11
             1
             32
             25
             35
             8
             17
             7
             9
             4
             2
             34
             10
             3
            "
        );

        let jolts = parse_jolts(input)?;
        let result = chain_all(&jolts);

        assert_eq!(result.0 * result.1, 22 * 10);
        Ok(())
    }

    #[test]
    fn test_simple3() -> Result<(), Box<dyn std::error::Error>> {
        let input = indoc!(
            "16
             10
             15
             5
             1
             11
             7
             19
             6
             12
             4
            "
        );

        let jolts = parse_jolts(input)?;
        let result = count_possible_chains(&jolts)?;
        assert_eq!(result, 8);

        Ok(())
    }

    #[test]
    fn test_simple4() -> Result<(), Box<dyn std::error::Error>> {
        let input = indoc!(
            "28
             33
             18
             42
             31
             14
             46
             20
             48
             47
             24
             23
             49
             45
             19
             38
             39
             11
             1
             32
             25
             35
             8
             17
             7
             9
             4
             2
             34
             10
             3
            "
        );

        let jolts = parse_jolts(input)?;
        let result = count_possible_chains(&jolts)?;

        assert_eq!(result, 19208);
        Ok(())
    }
}