aoc-2023-day-17.rs

use std::{fs::File, io::Read, path::Path};

use pathfinding::prelude::dijkstra;

const PART1_MAX_STRAIGHT: usize = 3;
const PART2_MIN_STRAIGHT: usize = 4;
const PART2_MAX_STRAIGHT: usize = 10;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum Dir {
  Up,
  Down,
  Left,
  Right,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct Coord(usize, usize);

impl Coord {
  fn new(x: usize, y: usize) -> Self {
    Coord(x, y)
  }

  fn go(&self, dir: Dir, map: &Map) -> Option<Self> {
    let (x, y) = (self.0 as i64, self.1 as i64);
    let (new_x, new_y) = match dir {
      Dir::Up => (x, y - 1),
      Dir::Down => (x, y + 1),
      Dir::Left => (x - 1, y),
      Dir::Right => (x + 1, y),
    };

    if !map.in_bound(new_x, new_y) {
      return None;
    }

    Some(Coord::new(new_x as usize, new_y as usize))
  }
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
struct Node {
  coord: Coord,
  history: (Dir, usize),
}

impl Node {
  fn new(coord: Coord, history: (Dir, usize)) -> Self {
    Node { coord, history }
  }

  fn cost(&self, map: &Map) -> usize {
    map.get(self.coord.0, self.coord.1)
  }

  fn neighbors_part1(&self, map: &Map) -> Vec<(Self, usize)> {
    let mut neighbors = Vec::new();

    for dir in [Dir::Up, Dir::Down, Dir::Left, Dir::Right] {
      let Some(new_coord) = self.coord.go(dir, map) else {
        continue;
      };

      let (last_dir, last_count) = self.history;

      // can't turn 180 degree
      if last_dir == Dir::Up && dir == Dir::Down
        || last_dir == Dir::Down && dir == Dir::Up
        || last_dir == Dir::Left && dir == Dir::Right
        || last_dir == Dir::Right && dir == Dir::Left
      {
        continue;
      }

      let new_history = if last_dir == dir && last_count < PART1_MAX_STRAIGHT {
        (last_dir, last_count + 1)
      } else if last_dir != dir {
        (dir, 1)
      } else {
        continue;
      };
      let new_node = Node::new(new_coord, new_history);

      let cost = new_node.cost(map);

      neighbors.push((new_node, cost));
    }

    neighbors
  }

  fn neighbors_part2(&self, map: &Map) -> Vec<(Self, usize)> {
    let mut neighbors = Vec::new();

    for dir in [Dir::Up, Dir::Down, Dir::Left, Dir::Right] {
      let Some(new_coord) = self.coord.go(dir, map) else {
        continue;
      };

      let (last_dir, last_count) = self.history;

      // can't turn 180 degree
      if last_dir == Dir::Up && dir == Dir::Down
        || last_dir == Dir::Down && dir == Dir::Up
        || last_dir == Dir::Left && dir == Dir::Right
        || last_dir == Dir::Right && dir == Dir::Left
      {
        continue;
      }

      let new_history = if last_dir == dir && last_count < PART2_MAX_STRAIGHT {
        (last_dir, last_count + 1)
      } else if last_dir != dir && last_count >= PART2_MIN_STRAIGHT {
        (dir, 1)
      } else {
        continue;
      };

      let new_node = Node::new(new_coord, new_history);
      let cost = new_node.cost(map);

      neighbors.push((new_node, cost));
    }

    neighbors
  }

  fn is_terminal(&self, map: &Map) -> bool {
    self.coord.0 == map.width - 1 && self.coord.1 == map.height - 1
  }
}

struct Map {
  width: usize,
  height: usize,
  tile: Vec<usize>,
}

impl Map {
  fn load(f: &Path) -> Self {
    let mut map = Map {
      width: 0,
      height: 0,
      tile: Vec::new(),
    };

    let mut file = File::open(f).unwrap();
    let mut contents = String::new();
    file.read_to_string(&mut contents).unwrap();

    map.height = contents.lines().count();
    map.width = contents.lines().next().unwrap().chars().count();

    let lines = contents.lines();

    for line in lines {
      for tile in line.trim_end().chars() {
        map.tile.push(tile.to_digit(10).unwrap() as usize);
      }
    }

    map
  }

  fn get(&self, x: usize, y: usize) -> usize {
    self.tile[y * self.width + x]
  }

  fn in_bound(&self, x: i64, y: i64) -> bool {
    x >= 0 && y >= 0 && x < self.width as i64 && y < self.height as i64
  }
}

fn main() {
  let map = Map::load(Path::new("./map.txt"));

  let out = dijkstra(
    &Node::new(Coord::new(0, 0), (Dir::Right, 0)),
    |node| node.neighbors_part2(&map),
    |node| node.is_terminal(&map),
  );

  if let Some((_, cost)) = out {
    println!("cost: {cost}");
  } else {
    println!("no path found");
  }
}