aoc2024/day12/src/main.rs

use std::collections::{HashMap, HashSet, VecDeque};

fn main() {
    let input = std::fs::read_to_string("input").unwrap();
    let grid = Garden::new(&input);
    println!("{}", pt1(&grid));
}

fn pt1(grid: &Garden) -> usize {
    let mut cost = 0;
    let regions = regions(grid);
    println!("there are {} regions", regions.len());
    for (_root, region) in regions {
        let a = region.len();
        let mut p = 0;
        for plot in region {
            p += plot.p_val;
        }
        cost += a * p;
    }

    cost
}

fn regions(grid: &Garden) -> Regions {
    let mut out = Regions::new();
    let mut knowns = HashSet::new();

    for row in grid.rows.iter() {
        for plot in row.iter() {
            let root = plot.loc;
            let mut q = VecDeque::new();
            if !knowns.contains(&root) {
                q.push_back(root);
            } else {
                continue;
            }
            let mut region = HashSet::new();
            while let Some(current) = q.pop_front() {
                knowns.insert(current);
                let nexts = grid.neighbors(&current);
                for next in nexts
                    .iter()
                    .map(|n| grid.get(n).unwrap())
                    .filter(|n| !region.contains(n))
                {
                    q.push_back(next.loc);
                }
                region.insert(grid.get(&current).unwrap());
            }
            let _ = out.entry(root).or_insert(region);
        }
    }
    out
}

type Regions<'g> = HashMap<Loc, HashSet<&'g Plot>>;

#[derive(Debug, Clone)]
struct Garden {
    rows: Vec<Vec<Plot>>,
}

impl Garden {
    fn new(input: &str) -> Self {
        let mut rows = Vec::new();
        for (row, line) in input.lines().enumerate() {
            let row = line
                .chars()
                .enumerate()
                .map(|(col, c)| Plot {
                    loc: Loc { row, col },
                    p_val: 0,
                    plant: c,
                })
                .collect();
            rows.push(row);
        }

        let mut g = Self { rows };
        let rows = g.rows.len();
        let cols = g.rows[0].len();
        for row in 0..rows {
            for col in 0..cols {
                let mut plot = g.rows[row][col];
                let mut p = 4;
                let loc = plot.loc;
                for _ in g.neighbors(&loc) {
                    p -= 1;
                }
                plot.p_val = p;
                g.rows[row][col] = plot;
            }
        }
        g
    }

    fn get(&self, loc: &Loc) -> Option<&Plot> {
        self.rows.get(loc.row).and_then(|row| row.get(loc.col))
    }

    fn next(&self, loc: &Loc) -> Vec<&Plot> {
        let Loc { row, col } = *loc;
        let mut out = Vec::new();

        // north
        {
            let row = row.wrapping_sub(1);
            if let Some(plot) = self.rows.get(row).and_then(|r| r.get(col)) {
                out.push(plot)
            }
        }

        // south
        {
            let row = row + 1;
            if let Some(plot) = self.rows.get(row).and_then(|r| r.get(col)) {
                out.push(plot);
            }
        }

        // east
        {
            let col = col + 1;
            if let Some(plot) = self.rows.get(row).and_then(|r| r.get(col)) {
                out.push(plot);
            }
        }

        // west
        {
            let col = col.wrapping_sub(1);
            if let Some(plot) = self.rows.get(row).and_then(|r| r.get(col)) {
                out.push(plot);
            }
        }
        out
    }

    // returns list of neighbor nodes in a dag where the edges are all one step in
    // distance
    fn neighbors(&self, loc: &Loc) -> Vec<Loc> {
        let plot = self.rows[loc.row][loc.col];
        self.next(loc)
            .iter()
            .filter_map(|t| {
                if t.plant == plot.plant {
                    Some(t.loc)
                } else {
                    None
                }
            })
            .collect()
    }
}

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

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
struct Plot {
    loc: Loc,
    p_val: usize,
    plant: char,
}

#[cfg(test)]
mod test {

    use super::*;

    static INPUT: &str = "RRRRIICCFF
RRRRIICCCF
VVRRRCCFFF
VVRCCCJFFF
VVVVCJJCFE
VVIVCCJJEE
VVIIICJJEE
MIIIIIJJEE
MIIISIJEEE
MMMISSJEEE";

    #[test]
    fn p1() {
        let g = Garden::new(INPUT);
        assert_eq!(1930, pt1(&g));
    }
}