18.rs

use std::time::Instant;
use std::collections::{HashMap, HashSet, VecDeque};
use itertools::Itertools;

const INPUT: [&str; 81] = [
  "#################################################################################",
  "#.............#...#...O.#.#...........#.#...#.........#.......#.....#.......#.Z.#",
  "#####.#######.#H#.#.###.#.#.#####.###.#.#.###.#####.#.#.#####.#.###.#.###.###.#.#",
  "#.....#.#...#.#.#...#.#.#.#.#.#...#.....#.#...#...#.#.#...#w..#...#...#.#.....#.#",
  "#.#####.#.#.#.#.#####.#.#.#B#.#.#########.#.###.#.#.#####.#.#####.#####.#######.#",
  "#.#...#...#.....#...#.#.#.....#.........#.#.#...#.#.....#.#...#..y..#...#.....#.#",
  "#.#.#.#.#########.#.#.#.###########.###.#.#.#.#######.###.###.#####.#.#.#.###.#.#",
  "#.#.#...#.........#.#.......#.....#...#.#.#.#.#.....#.....#.#.#...#...#.#.#...#.#",
  "#.###.###.#########.#######.#.###.#####.#.#.#.#.###.#######.#.#.#.#####.#.#.###.#",
  "#...#.#.....#x....#.#..f..#...#...#...#.#.#.#.....#.....#i..#...#.#...#.#.#.#...#",
  "#.#.#.#.#####.###.#.#A###.#####.###.#.#.#.#.#########.#.#.#.#####.#.#.###.#.#.#.#",
  "#.#.#.#.#...#.#.#.#...#.#.#...#.#.T.#...#.#...#.....#.#...#.#.#..e#.#.....#.#.#.#",
  "###.###.#X#.#.#.#.#.###.#.###.#.###.###.#.###.#.###.#.#####.#.#.###.#######.#.#.#",
  "#...#d..#.#...#.#.#.....#...#.#...#...#.#...#.#.#.....#.......#.#.....#.#...#.#.#",
  "#.###.###.#####R#.#####.###.#.###.###.#.#.#.#.#.###########.###.#.###.#.#.#####.#",
  "#...#.#.#...#.......#...#...#...#.....#.#.#...#.....#.....#.#...#...#.#...#...#.#",
  "###E#.#.###.#######.#####.###.#.#######.#.#####.###N#.###.###.###.###.#.###.#.#.#",
  "#.#...#...#..c..#...#r..F.#...#.....#.#.#...#.#.#.#.#...#.#...#...#...#.....#.#.#",
  "#.#####.#######.###.#.#######.#####.#.#.###.#.#.#.#.###.#.#.###.###.#########.#.#",
  "#z....#.......#...#...#.....#s#.......#.#.#.#.....#...#.#...#v..#.#...#.....#...#",
  "#.#.###.#.#######.#####G###.#.#######.#.#.#.#####.#.###.#######.#.###.#.#######.#",
  "#.#.#...#.......#...#...#.#.#...#...#.#.#.#.....#.#...#.D.....#.....#.#.......#.#",
  "###.#.#########.###C#####.#.#.#.#.#.###.#.#####.#####.###############.#.#####.#.#",
  "#...#...#.....#...#....g....#.#.#.#.....#.....#.....#...........#.....#.....#...#",
  "#.#####.#.###.#.#############.#.#.#######.#########.#.#######.#.#.#########.#####",
  "#.....#...#...#.#.........#...#.#.#.....#.........#.#.#.....#.#...#.......#.#..u#",
  "#Q#.#######.###.#.#.#.#####.###.#.#####.#.#######.#.###.###.#.#####.###.#.#.#.###",
  "#.#.......#.#...#.#.#.#...#.#.#...#...#.#.#.....#.#.#...#.#.#...#.....#.#.#.#...#",
  "#.#####.###.#.###.#.###.#.#.#.#####.#.#.#.#.###.#.#.#.###.#.#.#.#######.#.#.#.#.#",
  "#...#...#...#.....#.....#...#.......#...#.#.#.#.K.#...#.#...#.#.#.......#.#.#.#.#",
  "###.#.###.#####################.###.#####.#.#.#.#######.#.#####.#.#####.###.#.#.#",
  "#...#.....#...#.....#.........#...#.#...#.#.#...#.......#...#...#...#...#...#.#.#",
  "#.#########.###.###.#.###.###.###.#.#.#.###.#####.#.###.###.#.###.#.#.###.#####.#",
  "#.#.....#.........#...#...#.#.#.#.#.#.#.#...#.....#.#l..#.#.#...#.#.#...#...#...#",
  "#.###.#.###.###########.###.#.#.#.#.#.#.#.###.#####.#.###.#.###.#.#.#.#####.#.#.#",
  "#.#...#...#...#...#.....#...#.#...#...#.#...#.#.#...#.....#.....#.#.#.#...#.#.#.#",
  "#.#.#####.#.###.#.#.#####.###.#########.#.#.#.#.#.#################.###.#.#P#.#.#",
  "#.#.#.....#.#...#.#...#.#...#.#.....#...#.#.#.#.#.#.................#...#...#.#.#",
  "#.#.#.#######.###.###.#.#.#.#.#.###.#.#.#.#.#.#.#.#.#########.#######.#######.#.#",
  "#.U.#p........#....j..#...#.....#.....#...#.....#...........#...........J.....#.#",
  "#######################################.@.#######################################",
  "#...#.....#...#.......#.........#.........#.....#.....#.....#.....#.............#",
  "#.#.#.###.###.#.#L###.#.#####.###.#.###.#.#.#.###.#.###.#.#.#.#.#.#.#########.#.#",
  "#.#...#.....#...#...#.#.#...#.....#...#.#...#.....#.....#.#.#.#.#.#.#.......#q#.#",
  "#.#.#######.#.#####.#.###.#.#########.#.#.###############.###.#.#.###.#####.###.#",
  "#.#.#.....#.#...#...#.....#.#.........#.#.#.....#.....#...#...#.#.....#...#...#.#",
  "#.#.#.###.#.#####.#########.#.#########.#.#.#.#.#.###.#.###.###.#.#####.#.###.#.#",
  "#.#.#...#.#.....#.....#...#...#.......#.#.#.#.#...#...#.#.....#.#.#...#.#...#...#",
  "#.#####.#.#####.#####.#.#######.###.###.#.#.#.#####.###.#.#####.#.#.#.#.###.###.#",
  "#.....#.#...#...#...#.#...........#...#.#k#.#...#.....#.#.#.....#.#.#.#...#...#.#",
  "#####.#.###.###.#.#.#.#.#############.#.#.#.###.#.#####.#.#.#######.#.###.#.###.#",
  "#.....#...#...#...#...#.#...........#...#.#...#.#...#...#.#...#...#.#...#.#.....#",
  "#.#######.###.#####.#####.#########.#.###.#####.###.#.###.###.#.#.#.###.#.#######",
  "#.......#...#.....#.#.....#.......#.#...#.......#.#.#.#...#.#...#.....#.#.#.....#",
  "#.###.#####.#####.###.#####.#.#.###.#############.#.#.#.###.###########.#.#.###.#",
  "#...#.#.........#...#.#.....#.#.#...#...#...#.....#.#...#...#.......#...#.#.#m..#",
  "#.###.#.###########.#.#######.###.###.#.#.#.#.#.###.#####.#.#######.#.###.#.#.#.#",
  "#.#...#...#.#.......#...#...#.W...#...#.#.#.#.#.....#.....#...#.....#...#.#.#.#.#",
  "###.#.###.#.#.#########.#.#.#####.###.#.#.###.#######.#######.#.###.###.#.###.#.#",
  "#...#...#...#...#.....#...#.....#...#.#.#.#...#.........#.......#.#.#...#.....#.#",
  "#.#########.###.#.#.###########.###.#.#.#.#.###.#######.#.#######.#.#.#########.#",
  "#.........#.#...#.#.#....a......#.#...#.#.#...#.#.......#.#...#.....#.#...#.....#",
  "#.#######.#.#.###.#.#############.#####.#.###.#.#.#######.#.###.#####.###.#.#####",
  "#.#.....#...#.#...#.....#...#.......#h..#.....#.#.#.......#.#...#...#...#.#.#...#",
  "#.###.#.#####.#.#####.#.#.#.#.#####.#.###.#######.#######.#.#####.#.###.#.#.#.#.#",
  "#...#.#...#...#.#...#.#.#.#..o#...#...#.#.......#.......#.#.....#.#.#...#.....#.#",
  "###.#####.#.#.#.###.#.###.#######.#####.#######.#.#####.#######.#.#.#.###########",
  "#.#...#...#.#.#.....#.....#.#.......#...#...#...#.....#.......#.#.#.#...#.......#",
  "#.###.#.###.#.#####.#######.#.#I###.###.#.#.#.###########.###.#.#.#.###.#.#####.#",
  "#...#...#...#.....#.#.#.....#.#...#t..#.#.#...#.........#.#.#.#...#...#...#.....#",
  "#.#####.#.#########.#.#.###.#.###.###.#.#.#####.#######.#.#.#.#######.#####.#####",
  "#...M...#.......#...#...#.#...#.#...#...#.......#.....#.#.#.........#.....#.#...#",
  "#.#############.#.###.###.#####.###.#############.#####.#.#.#######.#####.#.###.#",
  "#.#.....Y...#.#.#.#.#.#..b..#...#.#.#...#...#.........#.#.#.#...#.#.#...#...#...#",
  "#.#.#######.#.#.#.#.#.###.#.#.#.#.#.#.#.#.#.#.#######.#.#.###.#.#.#.#.#.#####.#.#",
  "#...#n....#.#.#...#.......#...#...#...#.#.#.#.#.....#...#.....#.#...#.#.......#.#",
  "#####.###.#.#.###################.#####.#.###.#.###.#####V#####.#.###.###.#####.#",
  "#.....#...#...#.....#...#...#.....#.....#.....#...#...........#.#...#.#...#.....#",
  "#.###########.###.#.#.#.#.#.#######.###.#.#####################.###.###.###.#####",
  "#.................#...#...#.........#.S.#.......................#.......#.......#",
  "#################################################################################",
];

type Pos = (usize, usize);
type Graph = HashMap<Pos, Vec<(Pos,usize)>>;
type Map = [Vec<char>];
type Path = (usize, usize, usize);

// Remove all nodes that are empty spaces with just one neighbour, i.e dead-ends
fn remove_dead_ends(nodes: &mut HashSet<Pos>, map: &Map) {
  loop {
    let to_remove = nodes.iter()
      .filter(|&&(i,j)| map[i][j] == '.')
      .filter(|&&(i,j)| {
        let mut n = 0;
        if nodes.contains(&(i-1,j)) { n += 1; }
        if nodes.contains(&(i+1,j)) { n += 1; }
        if nodes.contains(&(i,j-1)) { n += 1; }
        if nodes.contains(&(i,j+1)) { n += 1; }
        n == 1
      })
      .cloned()
      .collect_vec();
    if to_remove.is_empty() { return; }
    for pos in &to_remove { nodes.remove(pos); }
  }
}

// Contract nodes that are empty spaces and only functions as bridges between other nodes.
// This does not "change" the graph, it simply removes all uninteresting nodes,
// leaving only 3/4-way intersections and doors/keys as nodes in the graph.
fn contract_corridors(g: &mut Graph, map: &Map) {
  let to_contract = g.iter()
    .filter(|&((x,y), _)| map[*x][*y] == '.')
    .filter(|&(_, n)| n.len() == 2)
    .map(|(k,_)| *k)
    .collect_vec();
  for p in &to_contract {
    let (p1, d1) = g[p][0];
    let (p2, d2) = g[p][1];
    g.entry(p1).and_modify(|n| {
      let index = n.iter().position(|(pos,_)| pos == p).unwrap();
      n[index] = (p2,d1+d2);
    });
    g.entry(p2).and_modify(|n| {
      let index = n.iter().position(|(pos,_)| pos == p).unwrap();
      n[index] = (p1,d1+d2);
    });
    g.remove(&p);
  }
}

// This process reduces the huge maze above to just 79 nodes.
fn to_minimized_graph(map: &Map) -> Graph {
  let mut nodes = (0..map.len()).cartesian_product(0..map[0].len())
    .filter(|&(x,y)| map[x][y] != '#')
    .collect::<HashSet<_>>();
  remove_dead_ends(&mut nodes, &map);

  let mut g = nodes.iter().map(|&(i,j)| {
    let n = [(i-1,j), (i+1,j), (i,j-1), (i,j+1)].iter()
      .filter(|p| nodes.contains(p))
      .map(|&p| (p,1))
      .collect();
    ((i,j), n)
  }).collect();
  contract_corridors(&mut g, &map);
  g
}

// Store keys as bits instead so we can speed up computations via bitwise operations
fn char_to_bit(c: char, offset: u8) -> usize {
  match c {
    '@' => 1 << 26,
    '%' => 1 << 27,
    '!' => 1 << 28,
    '€' => 1 << 29,
    _   => 1 << (c as u8 - offset),
  }
}

fn bfs(g: &Graph, map: &Map, start: Pos, end: Pos) -> Option<Path> {
  let mut queue = VecDeque::new();
  let mut visited = HashSet::new();
  queue.push_back(vec![(start,0)]);

  let path = loop {
    let path = match queue.pop_front() {
      Some(path) => path,
      None       => return None,
    };
    let (u, _) = *path.last().unwrap();
    if u == end { break path; }

    visited.insert(u);
    g[&u].iter()
      .filter(|&(v,_)| !visited.contains(v))
      .for_each(|&(v,d)| {
        let mut new_path = path.clone();
        new_path.push((v,d));
        queue.push_back(new_path);
      })
  };

  let doors = path.iter().map(|&((x,y),_)| map[x][y]).filter(|c| c.is_ascii_uppercase()).fold(0, |mask, c| mask | char_to_bit(c, b'A'));
  let keys  = path.iter().map(|&((x,y),_)| map[x][y]).filter(|c| c.is_ascii_lowercase()).fold(0, |mask, c| mask | char_to_bit(c, b'a'));
  let len   = path.iter().map(|(_,d)| d).sum();
  Some((doors, keys, len))
}

// Compute the length of path between all key pairs,
// what keys are picked up along the way,
// and what doors are required to get there.
fn compute_paths(g: &Graph, map: &Map) -> HashMap<usize, HashMap<usize, Path>> {
  let keys = g.keys()
    .filter(|&&(x,y)| map[x][y].is_ascii_lowercase() || ['@','%','!','€'].contains(&map[x][y]))
    .cloned()
    .collect_vec();
  let mut paths = HashMap::new();
  for i in 0..keys.len() {
    let (x1,y1) = keys[i];
    let key1 = char_to_bit(map[x1][y1], b'a');
    for j in i..keys.len() {
      let (x2,y2) = keys[j];
      let key2 = char_to_bit(map[x2][y2], b'a');
      if let Some(path) = bfs(&g, &map, keys[i], keys[j]) {
        paths.entry(key1).or_insert(HashMap::new()).insert(key2, path);
        paths.entry(key2).or_insert(HashMap::new()).insert(key1, path);
      }
    }
  }
  paths
}

fn shortest_path(
  g: &Graph,
  paths: &HashMap<usize, HashMap<usize,Path>>,
  cache: &mut HashMap<(usize,usize), usize>,
  all_keys: usize,
  keys: usize,
  curr: usize,
) -> usize {
  if keys == all_keys { return 0; }
  if let Some(&d) = cache.get(&(keys,curr)) { return d; }
  let d = paths[&curr].iter()
    .filter(|&(key,_)| key & keys == 0)
    .filter(|(_,(doors,_,_))| keys & doors == *doors)
    .map(|(&key, (_, new_keys, len))|
      len + shortest_path(g, paths, cache, all_keys, keys | new_keys, key)
    )
    .min()
    .unwrap();
  cache.insert((keys,curr), d);
  d
}

fn part_one(map: &Map) -> usize {
  let g = to_minimized_graph(&map);
  let paths = compute_paths(&g, &map);
  let start = char_to_bit('@',0);
  let all_keys = (1 << 27) - 1;
  shortest_path(&g, &paths, &mut HashMap::new(), all_keys, start, start)
}

fn part_two(map: &mut Map) -> usize {
  map[39][39] = '@'; map[40][39] = '#'; map[41][39] = '!';
  map[39][40] = '#'; map[40][40] = '#'; map[41][40] = '#';
  map[39][41] = '%'; map[40][41] = '#'; map[41][41] = '€';

  let g = to_minimized_graph(&map);
  let mut paths = compute_paths(&g, &map);

  // Any doors this robot cannot open will be opened by
  // another robot at some point so we can just remove
  // those doors before calculating the shortest path.
  ['@','%','!','€'].iter().map(|&c| {
    let k = char_to_bit(c,0);
    let keys = paths[&k].keys().cloned().collect_vec();
    let reachable_keys = keys.iter().fold(k, |ks, k| ks | k);
    for k in keys {
      paths.entry(k).and_modify(|p| {
        for (_, (doors,_,_)) in p.iter_mut() {
          *doors &= reachable_keys;
        }
      });
    }
    shortest_path(&g, &paths, &mut HashMap::new(), reachable_keys, k, k)
  }).sum()
}

fn main() {
  let now = Instant::now();
  let mut map = INPUT.iter().map(|s| s.chars().collect_vec()).collect_vec();
  let part_one = part_one(&map);
  let part_two = part_two(&mut map);
  println!("Part one: {}", part_one);
  println!("Part two: {}", part_two);
  println!("Time: {}ms", now.elapsed().as_millis());
}