advent-of-code/2023/src/day3.rs

use std::collections::HashSet;

const INPUT: &str = include_str!("../data/day3.txt");

pub fn day3a() -> String {
    format!("{}", sum_parts(INPUT.to_owned()))
}

pub fn day3b() -> String {
    format!(
        "{}",
        find_gears(INPUT.to_owned())
            .into_iter()
            .map(|gear| gear.ratio())
            .fold(0, |acc, val| acc + val)
    )
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct Addr {
    row: usize,
    col: usize,
}

impl Addr {
    fn halo(&self) -> impl Iterator<Item = Addr> {
        let row = self.row as i32;
        let col = self.col as i32;
        vec![
            (row - 1, col - 1),
            (row, col - 1),
            (row + 1, col - 1),
            //
            (row - 1, col),
            (row, col),
            (row + 1, col),
            //
            (row - 1, col + 1),
            (row, col + 1),
            (row + 1, col + 1),
        ]
        .into_iter()
        .filter(|(row, col)| *row >= 0 && *col >= 0)
        .map(|(row, col)| Addr {
            row: row as usize,
            col: col as usize,
        })
    }
}

#[derive(Clone, Debug, PartialEq)]
struct Part {
    value: u32,
    row: usize,
    start: usize,
    end: usize,
}

impl Part {
    fn span<'a>(&'a self) -> impl Iterator<Item = Addr> + 'a {
        (self.start..self.end + 1).map(|addr| Addr {
            row: self.row,
            col: addr,
        })
    }

    fn halo<'a>(&'a self, max_row: usize, max_col: usize) -> impl IntoIterator<Item = Addr> {
        println!("{:?}", self);
        let addresses = self
            .span()
            .map(|addr| addr.halo())
            .flatten()
            .filter(|Addr { row, col }| {
                *row as i32 >= 0
                    && *row as i32 <= max_row as i32
                    && *col as i32 >= 0
                    && *col as i32 <= max_col as i32
            })
            .fold(HashSet::new(), |mut acc, val| {
                let _ = acc.insert(val);
                acc
            });

        println!("{:?}", addresses);
        println!("{:?}", self.span().collect::<HashSet<Addr>>());
        println!("{:?}", addresses.difference(&self.span().collect()));

        addresses
            .difference(&self.span().collect())
            .map(|addr| addr.clone())
            .collect::<HashSet<Addr>>()
    }
}

#[derive(Clone, Debug, PartialEq)]
struct Gear(u32, u32);

impl Gear {
    fn ratio(&self) -> u32 {
        self.0 * self.1
    }
}

fn sum_parts(schematic: String) -> u32 {
    parts(schematic)
        .into_iter()
        .fold(0, |acc, val| acc + val.value)
}

fn parts(schematic: String) -> Vec<Part> {
    let sc = schematic
        .lines()
        .map(|l| l.to_owned())
        .collect::<Vec<String>>();
    let parts = find_numbers(schematic.clone());
    parts
        .into_iter()
        .filter(|part| is_part(&sc, part.clone()))
        .collect::<Vec<Part>>()
}

fn find_gears(schematic: String) -> Vec<Gear> {
    let sc = schematic
        .lines()
        .map(|l| l.to_owned())
        .collect::<Vec<String>>();
    let max_row = sc.len();
    let max_col = sc[0].len();
    let parts = parts(schematic);
    let splats = find_splats(sc)
        .into_iter()
        .map(|addr| {
            Part {
                value: 0,
                row: addr.row,
                start: addr.col,
                end: addr.col,
            }
            .halo(max_row - 1, max_col - 1)
            .into_iter()
            .collect::<HashSet<Addr>>()
        })
        .collect::<Vec<HashSet<Addr>>>();

    let mut gears = vec![];
    for splat in splats {
        let overlap = parts
            .iter()
            .filter(|part| part_within_halo(splat.clone(), (*part).clone()))
            .map(|part| part.clone())
            .collect::<Vec<Part>>();
        if overlap.len() == 2 {
            gears.push(Gear(overlap[0].value, overlap[1].value));
        }
    }

    gears
}

fn part_within_halo(halo1: impl IntoIterator<Item = Addr>, part: Part) -> bool {
    let halo1: HashSet<Addr> = halo1.into_iter().collect::<HashSet<Addr>>();
    let halo2: HashSet<Addr> = part.span().into_iter().collect::<HashSet<Addr>>();

    halo1.intersection(&halo2).count() > 0
}

fn is_part(schematic: &Vec<String>, part: Part) -> bool {
    part.halo(schematic.len() - 1, schematic[0].len() - 1)
        .into_iter()
        .any(|addr| {
            let cell = schematic[addr.row].chars().nth(addr.col).unwrap();
            !cell.is_digit(10) && cell != '.'
        })
}

fn find_numbers(schematic: String) -> Vec<Part> {
    let mut result = vec![];
    for (idx, line) in schematic.lines().enumerate() {
        let mut numbers = find_numbers_in_line(line)
            .into_iter()
            .map(|(value, start, len)| Part {
                value: value,
                row: idx,
                start,
                end: start + len - 1,
            })
            .collect::<Vec<Part>>();
        result.append(&mut numbers);
    }
    result
}

fn find_numbers_in_line(line: &str) -> Vec<(u32, usize, usize)> {
    let mut numbers = vec![];

    let mut start_idx: usize = 0;
    let mut number: Option<String> = None;

    for idx in 0..line.len() {
        match (number.take(), line.chars().nth(idx).unwrap()) {
            (None, c) if c.is_digit(10) => {
                start_idx = idx;
                number = Some(String::from(c));
            }
            (Some(mut num), c) if c.is_digit(10) => {
                num.push(c);
                number = Some(num);
            }
            (Some(num), _) => {
                numbers.push((num, start_idx));
                number = None;
            }
            (_, _) => {}
        }
    }

    if let Some(num) = number {
        numbers.push((num, start_idx));
    }

    numbers
        .into_iter()
        .map(|(number, start_idx)| (number.parse::<u32>().unwrap(), start_idx, number.len()))
        .collect()
}

fn find_splats(schematic: Vec<String>) -> Vec<Addr> {
    schematic
        .into_iter()
        .enumerate()
        .map(|(row_idx, row)| {
            row.chars()
                .enumerate()
                .filter_map(|(col_idx, c)| {
                    if is_splat(c) {
                        Some(Addr {
                            row: row_idx,
                            col: col_idx,
                        })
                    } else {
                        None
                    }
                })
                .collect::<Vec<Addr>>()
        })
        .flatten()
        .collect::<Vec<Addr>>()
}

fn is_symbol(c: char) -> bool {
    !c.is_digit(10) && c != '.'
}

fn is_splat(c: char) -> bool {
    c == '*'
}

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

    const INPUT: &str = "467..114..
...*......
..35..633.
......#...
617*......
.....+.58.
..592.....
......755.
...$.*....
.664.598..";

    #[test]
    fn part_1() {
        assert_eq!(
            parts(INPUT.to_owned())
                .into_iter()
                .map(|part| part.value)
                .collect::<Vec<u32>>(),
            vec![467, 35, 633, 617, 592, 755, 664, 598]
        );

        assert_eq!(sum_parts(INPUT.to_owned()), 4361);
    }

    #[test]
    fn test_find_numbers() {
        println!("{:?}", find_numbers(INPUT.to_owned()));
        assert_eq!(
            find_numbers(INPUT.to_owned()),
            vec![
                Part {
                    value: 467,
                    row: 0,
                    start: 0,
                    end: 2
                },
                Part {
                    value: 114,
                    row: 0,
                    start: 5,
                    end: 7
                },
                Part {
                    value: 35,
                    row: 2,
                    start: 2,
                    end: 3
                },
                Part {
                    value: 633,
                    row: 2,
                    start: 6,
                    end: 8
                },
                Part {
                    value: 617,
                    row: 4,
                    start: 0,
                    end: 2
                },
                Part {
                    value: 58,
                    row: 5,
                    start: 7,
                    end: 8,
                },
                Part {
                    value: 592,
                    row: 6,
                    start: 2,
                    end: 4,
                },
                Part {
                    value: 755,
                    row: 7,
                    start: 6,
                    end: 8,
                },
                Part {
                    value: 664,
                    row: 9,
                    start: 1,
                    end: 3,
                },
                Part {
                    value: 598,
                    row: 9,
                    start: 5,
                    end: 7,
                }
            ]
        );
    }

    #[test]
    fn test_find_numbers_in_line() {
        let test_cases = vec![
            ("...*......", vec![]),
            (".664.598..", vec![(664, 1, 3), (598, 5, 3)]),
            ("..35..633.", vec![(35, 2, 2), (633, 6, 3)]),
            (
                "......883.....*163...782",
                vec![(883, 6, 3), (163, 15, 3), (782, 21, 3)],
            ),
        ];
        for (input, expected) in test_cases {
            assert_eq!(find_numbers_in_line(input), expected);
        }
    }

    #[test]
    fn test_halo() {
        let schematic = INPUT.lines().map(|l| l.to_owned()).collect::<Vec<String>>();
        let part = Part {
            value: 0,
            row: 0,
            start: 0,
            end: 2,
        };
        assert_eq!(
            part.halo(schematic.len() - 1, schematic[0].len() - 1)
                .into_iter()
                .collect::<HashSet<Addr>>(),
            vec![
                Addr { row: 1, col: 0 },
                Addr { row: 1, col: 1 },
                Addr { row: 1, col: 2 },
                Addr { row: 0, col: 3 },
                Addr { row: 1, col: 3 }
            ]
            .into_iter()
            .collect::<HashSet<Addr>>()
        );

        assert_eq!(
            Part {
                value: 0,
                row: 1,
                start: 1,
                end: 1
            }
            .halo(schematic.len() - 1, schematic[0].len() - 1)
            .into_iter()
            .collect::<HashSet<Addr>>(),
            vec![
                Addr { row: 0, col: 1 },
                Addr { row: 2, col: 1 },
                Addr { row: 0, col: 0 },
                Addr { row: 1, col: 0 },
                Addr { row: 2, col: 0 },
                Addr { row: 0, col: 2 },
                Addr { row: 1, col: 2 },
                Addr { row: 2, col: 2 }
            ]
            .into_iter()
            .collect::<HashSet<Addr>>()
        );
        assert_eq!(
            Part {
                value: 0,
                row: 6,
                start: 2,
                end: 4,
            }
            .halo(schematic.len() - 1, schematic[0].len() - 1)
            .into_iter()
            .collect::<HashSet<Addr>>(),
            vec![
                Addr { row: 5, col: 2 },
                Addr { row: 7, col: 2 },
                Addr { row: 5, col: 3 },
                Addr { row: 7, col: 3 },
                Addr { row: 5, col: 4 },
                Addr { row: 7, col: 4 },
                Addr { row: 5, col: 1 },
                Addr { row: 6, col: 1 },
                Addr { row: 7, col: 1 },
                Addr { row: 5, col: 5 },
                Addr { row: 6, col: 5 },
                Addr { row: 7, col: 5 },
            ]
            .into_iter()
            .collect::<HashSet<Addr>>()
        );
    }

    #[test]
    fn test_find_gears() {
        let gears = find_gears(INPUT.to_owned());
        assert_eq!(gears, vec![Gear(467, 35), Gear(755, 598)]);
        assert_eq!(
            gears.into_iter().map(|g| g.ratio()).collect::<Vec<u32>>(),
            vec![16345, 451490]
        );
    }
}