2022/src/bin/day13.rs

use std::{cmp::Ordering, fmt::Display, str::FromStr};

use aoc_2022::*;

type Input = Vec<Pair>;
type Output = usize;

#[derive(Debug, Clone, PartialEq, Eq)]
enum Packet {
    Integer(i32),
    List(Vec<Packet>),
}

impl Packet {
    fn is_integer(&self) -> bool {
        matches!(self, Packet::Integer(_))
    }
}

impl PartialOrd for Packet {
    fn partial_cmp(&self, right: &Self) -> Option<Ordering> {
        match (self, right) {
            (Packet::Integer(l), Packet::Integer(r)) => Some(l.cmp(r)),
            (Packet::List(l), Packet::List(r)) => {
                if let Some(x) = l
                    .iter()
                    .zip(r)
                    .map(|(l, r)| l.partial_cmp(r).unwrap())
                    .find(|x| x.is_ne())
                {
                    Some(x)
                } else {
                    Some(l.len().cmp(&r.len()))
                }
            }
            (l, r) => {
                if l.is_integer() {
                    Packet::List(vec![l.clone()]).partial_cmp(r)
                } else {
                    l.partial_cmp(&Packet::List(vec![r.clone()]))
                }
            }
        }
    }
}

impl Ord for Packet {
    fn cmp(&self, other: &Self) -> Ordering {
        self.partial_cmp(other).unwrap()
    }
}

impl Display for Packet {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Packet::Integer(x) => write!(f, "{x}"),
            Packet::List(lst) => write!(f, "[{}]", lst.iter().map(|p| format!("{p}")).join(",")),
        }
    }
}

impl FromStr for Packet {
    type Err = Report;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        fn parse(s: &str, mut i: usize) -> Result<(usize, Packet), Report> {
            // try to parse an integer
            let num = s
                .chars()
                .skip(i)
                .take_while(|c| c.is_ascii_digit())
                .join("");
            if !num.is_empty() {
                let integer = num.parse()?;
                return Ok((i + num.len(), Packet::Integer(integer)));
            }

            // we got a '[' and we should traverse the contents of the list
            i += 1;

            let mut packets = Vec::new();
            while i < s.len() && &s[i..i + 1] != "]" {
                if &s[i..i + 1] == "," {
                    // skipping commas
                    i += 1;
                }

                let (new_i, packet) = parse(s, i)?;

                i = new_i;
                packets.push(packet);
            }

            Ok((i + 1, Packet::List(packets)))
        }

        let parsed_packet = parse(s, 0)?;
        Ok(parsed_packet.1)
    }
}

#[derive(Debug)]
struct Pair(Packet, Packet);

impl Pair {
    fn right_order(&self) -> Ordering {
        let Pair(l, r) = self;
        l.cmp(r)
    }
}

impl FromStr for Pair {
    type Err = Report;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let (l, r) = s.split('\n').collect_tuple().unwrap();
        Ok(Pair(l.parse()?, r.parse()?))
    }
}

struct Day13;
impl Solution<Input, Output> for Day13 {
    fn parse_input<P: AsRef<Path>>(pathname: P) -> Input {
        file_to_string(pathname)
            .trim_end_matches('\n')
            .split("\n\n")
            .map(|p| p.parse().unwrap())
            .collect_vec()
    }

    fn part_1(input: &Input) -> Output {
        (1..)
            .zip(input.iter())
            .filter(|(_, p)| p.right_order().is_lt())
            .map(|(i, _)| i)
            .sum()
    }

    fn part_2(input: &Input) -> Output {
        let dividers = vec![
            Packet::List(vec![Packet::List(vec![Packet::Integer(2)])]),
            Packet::List(vec![Packet::List(vec![Packet::Integer(6)])]),
        ];

        let mut packets = input
            .iter()
            .flat_map(|Pair(l, r)| vec![l.clone(), r.clone()])
            .collect_vec();
        packets.append(&mut dividers.clone());

        packets.sort();

        // for p in &packets {
        //     debug!("{}", p);
        // }

        (1..)
            .zip(packets.iter())
            .filter(|(_, p)| dividers.contains(p))
            .map(|(i, _)| i)
            .product()
    }
}

fn main() -> Result<()> {
    // Day13::run("sample")
    Day13::main()
}

test_sample!(day_13, Day13, 13, 140);