Files
aoc2023/17/src/main.rs

482 lines
16 KiB
Rust

use colormap::ColorMap;
use std::collections::hash_map::RandomState;
use std::collections::{BinaryHeap, HashMap};
use std::fs::File;
use std::io::{BufRead, BufReader, Lines, Write};
use std::iter::repeat;
use std::time::Instant;
use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};
mod colormap;
const COLORMAP: &ColorMap = &colormap::COLORMAP_INFERNO;
// BOILERPLATE
type InputIter = Lines<BufReader<File>>;
fn get_input() -> InputIter {
let f = File::open("input").unwrap();
let br = BufReader::new(f);
br.lines()
}
fn main() {
let start = Instant::now();
let ans1 = problem1(get_input());
let duration = start.elapsed();
println!("Problem 1 solution: {} [{}s]", ans1, duration.as_secs_f64());
let start = Instant::now();
let ans2 = problem2(get_input());
let duration = start.elapsed();
println!("Problem 2 solution: {} [{}s]", ans2, duration.as_secs_f64());
}
// DATA
const UNPATH_CHAR: char = '█';
const UNVISITED_CHAR: char = ' ';
#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
enum Direction {
Left,
Right,
Up,
Down,
}
impl Direction {
const fn all() -> &'static [Self; 4] {
&[
Direction::Left,
Direction::Right,
Direction::Up,
Direction::Down,
]
}
const fn opposite(&self) -> Self {
match self {
Direction::Left => Direction::Right,
Direction::Right => Direction::Left,
Direction::Up => Direction::Down,
Direction::Down => Direction::Up,
}
}
}
impl From<&Direction> for char {
fn from(dir: &Direction) -> Self {
match dir {
Direction::Left => '←',
Direction::Right => '→',
Direction::Up => '↑',
Direction::Down => '↓',
}
}
}
struct CityMap {
map: Vec<Vec<u64>>,
}
impl CityMap {
fn offset_pos(&self, pos: (usize, usize), dir: &Direction) -> Option<(usize, usize)> {
match dir {
Direction::Left if pos.0 > 0 => Some((pos.0 - 1, pos.1)),
Direction::Right if pos.0 < self.map[0].len() - 1 => Some((pos.0 + 1, pos.1)),
Direction::Up if pos.1 > 0 => Some((pos.0, pos.1 - 1)),
Direction::Down if pos.1 < self.map.len() - 1 => Some((pos.0, pos.1 + 1)),
_ => None,
}
}
#[allow(dead_code)]
fn print(&self) -> Result<(), Box<dyn std::error::Error>> {
let cost_max = *self.map.iter().flat_map(|row| row.iter()).max().unwrap() as f64;
let mut so_lock = StandardStream::stdout(ColorChoice::Always);
for y in 0..self.map.len() {
for val in &self.map[y] {
so_lock.set_color(
ColorSpec::new()
.set_bg(Some(COLORMAP.apply(*val as f64 / cost_max)))
.set_fg(Some(Color::Black)),
)?;
so_lock.write_fmt(format_args!("{}", val))?;
}
so_lock.reset()?;
writeln!(so_lock)?;
}
Ok(())
}
}
type Position = (usize, usize);
struct WalkCost<'a> {
start: Position,
cost_from: Vec<Vec<HashMap<(Direction, usize), u64>>>,
map: &'a CityMap,
}
#[derive(Debug)]
struct Move {
new_pos: Position,
dir: &'static Direction,
consecutive: usize,
weight: u64,
}
impl PartialEq for Move {
fn eq(&self, other: &Self) -> bool {
self.weight == other.weight
}
}
impl Eq for Move {}
impl PartialOrd for Move {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl Ord for Move {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
std::cmp::Reverse(self.weight).cmp(&std::cmp::Reverse(other.weight))
}
}
impl<'a> WalkCost<'a> {
fn from_map(map: &'a CityMap, start: Position) -> Self {
Self {
map,
start,
cost_from: map
.map
.iter()
.map(|row| repeat(HashMap::new()).take(row.len()).collect())
.collect(),
}
}
fn compute(&mut self) {
let mut unvisited_next_moves: BinaryHeap<Move> = BinaryHeap::new();
let valid_start_moves: Vec<Move> = Direction::all()
.iter()
.filter_map(|dir| {
self.map.offset_pos(self.start, dir).and_then(|pos| {
Some(Move {
new_pos: pos,
dir,
consecutive: 1,
weight: self.map.map[pos.1][pos.0],
})
})
}) // valid positions
.collect();
for m in valid_start_moves {
unvisited_next_moves.push(m);
}
while let Some(cur_move) = unvisited_next_moves.pop() {
// we've been here already at a lower cost
// if cur_move.weight >= self.cost_to[cur_move.new_pos.1][cur_move.new_pos.0] {
// continue;
// }
if let Some(last_weight) = self.cost_from[cur_move.new_pos.1][cur_move.new_pos.0]
.get(&(*cur_move.dir, cur_move.consecutive))
{
if cur_move.weight < *last_weight {
self.cost_from[cur_move.new_pos.1][cur_move.new_pos.0]
.insert((*cur_move.dir, cur_move.consecutive), cur_move.weight);
} else {
continue;
} // visited before at lower cost }
} else {
self.cost_from[cur_move.new_pos.1][cur_move.new_pos.0]
.insert((*cur_move.dir, cur_move.consecutive), cur_move.weight);
}
// println!("state at {:?}: {:?}", cur_move, self.cost_from[cur_move.new_pos.1][cur_move.new_pos.0]);
let valid_moves = Direction::all().iter().filter_map(|dir| {
self.map
.offset_pos(cur_move.new_pos, dir)
.and_then(|new_pos| {
Some(Move {
new_pos,
dir,
consecutive: if cur_move.dir == dir {
cur_move.consecutive + 1
} else {
1
},
weight: cur_move.weight + self.map.map[new_pos.1][new_pos.0],
})
})
.filter(|m| m.consecutive != 4 && *m.dir != cur_move.dir.opposite())
.filter(|m| {
m.weight
< *self.cost_from[m.new_pos.1][m.new_pos.0]
.get(&(*m.dir, m.consecutive))
.unwrap_or(&u64::MAX)
})
}); // valid positions
// println!("valid moves with {:?}", cur_move);
for next_move in valid_moves {
// println!(" {:?}", next_move);
unvisited_next_moves.push(next_move);
}
}
}
}
struct WalkCost2<'a> {
start: Position,
cost_from: Vec<Vec<HashMap<(Direction, usize), u64>>>,
map: &'a CityMap,
}
impl<'a> WalkCost2<'a> {
fn from_map(map: &'a CityMap, start: Position) -> Self {
Self {
map,
start,
cost_from: map
.map
.iter()
.map(|row| repeat(HashMap::new()).take(row.len()).collect())
.collect(),
}
}
fn min_cost_at(&self, pos: Position) -> Option<&u64> {
self.cost_from[pos.1][pos.0].values().min()
}
fn compute(&mut self, to: Position) {
let mut unvisited_next_moves: BinaryHeap<Move> = BinaryHeap::new();
let valid_start_moves: Vec<Move> = Direction::all()
.iter()
.filter_map(|dir| {
self.map.offset_pos(self.start, dir).and_then(|pos| {
Some(Move {
new_pos: pos,
dir,
consecutive: 1,
weight: self.map.map[pos.1][pos.0],
})
})
}) // valid positions
.collect();
for m in valid_start_moves {
unvisited_next_moves.push(m);
}
while let Some(cur_move) = unvisited_next_moves.pop() {
// we've been here already at a lower cost
// if cur_move.weight >= self.cost_to[cur_move.new_pos.1][cur_move.new_pos.0] {
// continue;
// }
if let Some(last_weight) = self.cost_from[cur_move.new_pos.1][cur_move.new_pos.0]
.get(&(*cur_move.dir, cur_move.consecutive))
{
if cur_move.weight < *last_weight {
self.cost_from[cur_move.new_pos.1][cur_move.new_pos.0]
.insert((*cur_move.dir, cur_move.consecutive), cur_move.weight);
// println!("move {:?} inserted {:?}", cur_move, (*cur_move.dir, cur_move.consecutive));
} else {
continue;
} // visited before at lower cost }
} else {
// println!("move {:?} inserted {:?}", cur_move, (*cur_move.dir, cur_move.consecutive));
self.cost_from[cur_move.new_pos.1][cur_move.new_pos.0]
.insert((*cur_move.dir, cur_move.consecutive), cur_move.weight);
}
if cur_move.new_pos == to {
// println!("reached end pos {:?} via {:?}", to, cur_move);
continue;
}
let valid_moves = Direction::all().iter().filter_map(|dir| {
self.map
.offset_pos(cur_move.new_pos, dir)
.and_then(|new_pos| {
Some(Move {
new_pos,
dir,
consecutive: if cur_move.dir == dir {
cur_move.consecutive + 1
} else {
1
},
weight: cur_move.weight + self.map.map[new_pos.1][new_pos.0],
})
})
.filter(|m| m.new_pos != self.start)
.filter(|m| *m.dir != cur_move.dir.opposite())
.filter(|m| {
if m.dir == cur_move.dir {
m.consecutive < 11
} else {
cur_move.consecutive >= 4
}
})
.filter(|m| m.new_pos != to || m.consecutive >= 4)
.filter(|m| {
m.weight
< *self.cost_from[m.new_pos.1][m.new_pos.0]
.get(&(*m.dir, m.consecutive))
.unwrap_or(&u64::MAX)
})
}); // valid positions
// println!("valid moves with {:?}", cur_move);
for next_move in valid_moves {
// println!(" {:?}", next_move);
unvisited_next_moves.push(next_move);
}
}
}
fn shortest_path_to(&self, to: Position) -> Vec<(Position, Direction)> {
let mut path = Vec::new();
let mut cur_pos = to;
// start at the end, walk backwards
while cur_pos != self.start {
let (m, _val) = self.cost_from[cur_pos.1][cur_pos.0]
.iter()
.min_by(|a, b| a.1.cmp(b.1))
.unwrap();
path.push((cur_pos, m.0));
cur_pos = self.map.offset_pos(cur_pos, &m.0.opposite()).unwrap();
}
path
}
fn print_path(&self, to: Position) -> Result<(), Box<dyn std::error::Error>> {
let path = self.shortest_path_to(to);
let map: HashMap<_, _, RandomState> = HashMap::from_iter(path.into_iter());
let cost_max_of_min = *self
.cost_from
.iter()
.flat_map(|row| row.iter().filter_map(|cell| cell.values().min()))
.max()
.unwrap() as f64;
let mut so_lock = StandardStream::stdout(ColorChoice::Always);
for y in 0..self.cost_from.len() {
for x in 0..self.map.map[y].len() {
let mut color = ColorSpec::new();
let c = if let Some(to_dir) = map.get(&(x, y)) {
let normalized_cost =
*self.min_cost_at((x, y)).unwrap() as f64 / cost_max_of_min;
let bg_color = COLORMAP.apply(normalized_cost);
let fg_color = if let Color::Rgb(r, g, b) = bg_color {
Color::Rgb(255 - r, 255 - g, 255 - b)
} else {
Color::Black
};
color.set_fg(Some(fg_color)).set_bg(Some(bg_color)).bold();
to_dir.into()
} else {
if let Some(cost) = &self.min_cost_at((x, y)) {
color.set_fg(Some(COLORMAP.apply(**cost as f64 / cost_max_of_min)));
UNPATH_CHAR
} else {
color.set_fg(Some(Color::Black));
UNVISITED_CHAR
}
};
so_lock.set_color(&color)?;
let mut char_buf = [0u8; 4];
c.encode_utf8(&mut char_buf);
so_lock.write_all(&char_buf)?;
}
so_lock.reset()?;
writeln!(so_lock)?;
}
Ok(())
}
}
impl<T: BufRead> From<Lines<T>> for CityMap {
fn from(lines: Lines<T>) -> Self {
Self {
map: lines
.map(|l| l.unwrap().chars().map(|c| c as u64 - '0' as u64).collect())
.collect(),
}
}
}
// PROBLEM 1 solution
fn problem1<T: BufRead>(input: Lines<T>) -> u64 {
let map = CityMap::from(input);
let mut costs = WalkCost::from_map(&map, (0, 0));
costs.compute();
// println!("{}", costs);
// costs.print_shortest_path((costs.cost_to[0].len() - 1, costs.cost_to.len() - 1));
*costs.cost_from[costs.cost_from.len() - 1][costs.cost_from[0].len() - 1]
.values()
.min()
.unwrap()
}
// PROBLEM 2 solution
fn problem2<T: BufRead>(input: Lines<T>) -> u64 {
let map = CityMap::from(input);
// map.print().unwrap();
let mut costs = WalkCost2::from_map(&map, (0, 0));
costs.compute((map.map[0].len() - 1, map.map.len() - 1));
// println!("{}", costs);
costs
.print_path((costs.cost_from[0].len() - 1, costs.cost_from.len() - 1))
.unwrap();
*costs.cost_from[costs.cost_from.len() - 1][costs.cost_from[0].len() - 1]
.values()
.min()
.unwrap()
}
#[cfg(test)]
mod tests {
use crate::*;
use std::io::Cursor;
use test_case::test_case;
const EXAMPLE: &str = &"2413432311323
3215453535623
3255245654254
3446585845452
4546657867536
1438598798454
4457876987766
3637877979653
4654967986887
4564679986453
1224686865563
2546548887735
4322674655533";
const EXAMPLE2: &str = &"111111111111
999999999991
999999999991
999999999991
999999999991";
#[test]
fn problem1_example() {
let c = Cursor::new(EXAMPLE);
assert_eq!(problem1(c.lines()), 102);
}
#[test_case(EXAMPLE, 94)]
#[test_case(EXAMPLE2, 71)]
fn problem2_example(example: &str, expect: u64) {
let c = Cursor::new(example);
assert_eq!(problem2(c.lines()), expect);
}
}