main.rs

use std::fs::File;
use std::io::{BufRead, BufReader, Result};
use std::collections::{VecDeque, HashMap, HashSet};

type Pos = (i32, i32);

#[derive(Debug, Clone)]
enum Tile {
    Free, Key(usize), Door(usize),
}

#[derive(Debug)]
struct Map {
    tiles: HashMap<Pos, Tile>,
    keys: u32,
    pos: Pos,
}

impl Map {
    fn from_file(path: &str) -> Result<Map> {
        let f = File::open(path)?;
        let f = BufReader::new(f);

        let (mut tiles, mut keys) = (HashMap::new(), 0u32);

        let mut pos: Pos = (-1, -1);

        for (y, line) in f.lines().enumerate() {
            for (x, tile) in line.unwrap().chars().enumerate() {
                let p = (x as i32, y as i32);
                match tile {
                    'a'..='z' => {
                        let key = tile as usize - 'a' as usize;
                        tiles.insert(p, Tile::Key(key));
                        keys = keys | (1 << key) as u32;
                    },
                    'A'..='Z' => {
                        tiles.insert(p, Tile::Door(tile as usize - 'A' as usize));
                    },
                    '.' => {
                        tiles.insert(p, Tile::Free);
                    },
                    '@' => {
                        tiles.insert(p, Tile::Free);
                        pos = p;
                    },
                    _ => (),
                }
            }
        }

        Ok(Map {
            tiles: tiles,
            keys: keys,
            pos: pos,
        })
    }

    fn find_keys(map: &HashMap<Pos, Tile>, seen: &mut HashSet<Pos>, pos: Pos, keys: u32) -> u32 {
        seen.insert(pos);
        
        let mut found = match map.get(&pos) {
            Some(Tile::Key(key)) => (1 << key) as u32,
            _ => 0u32,
        };

        for (dx, dy) in &DIRS {
            let next_pos = (pos.0 + dx, pos.1 + dy);
            if !seen.contains(&next_pos) && map.contains_key(&next_pos) {
                found = found | Map::find_keys(map, seen, next_pos, keys);
            }
        }

        keys | found
    }

    fn partition(&self) -> Vec<Map> {
        let (cx, cy) = self.pos;

        let mut tiles = self.tiles.clone();

        tiles.remove(&self.pos);
        
        for (dx, dy) in &DIRS {
            tiles.remove(&(cx + dx, cy + dy));
        }

        let mut maps = Vec::new();

        for (dx, dy) in &[(-1, -1), (1, 1), (1, -1), (-1, 1)] {
            let part_pos = (cx + dx, cy + dy);
            tiles.insert(part_pos, Tile::Free);

            maps.push(Map {
                tiles: tiles.clone(),
                pos: part_pos,
                keys: Map::find_keys(&tiles, &mut HashSet::new(), part_pos, 0),
            });
        }
        
        maps
    }
}

#[derive(Hash, Eq, PartialEq, Copy, Clone)]
struct State {
    missing_keys: u32,
    pos: Pos,
}

const DIRS: [Pos; 4] = [(-1, 0), (1, 0), (0, 1), (0, -1)];

impl State {
    fn seed(map: &Map) -> State {
        State {
            missing_keys: map.keys,
            pos: map.pos,
        }
    }

    fn next(&self, map: &Map) -> Vec<State> {
        DIRS.into_iter()
            .filter_map(|(dx, dy)| {
                let next_pos = (self.pos.0 + dx, self.pos.1 + dy);
                
                match map.tiles.get(&next_pos) {
                    Some(Tile::Door(key)) => match (self.missing_keys >> key) & 1 {
                        0 => Some(State{missing_keys: self.missing_keys, pos: next_pos}),
                        _ => None,
                    },
                    Some(Tile::Key(key)) => Some(State{
                                                    missing_keys: match (self.missing_keys >> key) & 1 {
                                                        1 => self.missing_keys ^ (1 << key) as u32, 
                                                        _ => self.missing_keys,
                                                    },
                                                    pos: next_pos
                                                }),
                    Some(Tile::Free) => Some(State{missing_keys: self.missing_keys, pos: next_pos}),
                    _ => None,
                }
            }).collect()
    }
}

fn count_steps(map: &Map) -> Option<usize> {
    let steps_to: &mut HashMap<State, usize> = &mut HashMap::new();
    
    let seed = State::seed(map);

    let mut queue: VecDeque<State> = VecDeque::new();

    queue.push_back(seed);

    steps_to.insert(seed, 0);

    loop {
        match queue.pop_front() {
            Some(state) => match steps_to.get(&state) {
                Some(steps) => {
                    let steps_current = *steps;
                    
                    if state.missing_keys == 0 {
                        return Some(steps_current);
                    }
                    
                    for next_state in state.next(map) {
                        if !steps_to.contains_key(&next_state) {
                            steps_to.insert(next_state, steps_current + 1);

                            queue.push_back(next_state);
                        }
                    }
                },
                _ => break,
            }
            _ => break,
        }
    }

    None
}

fn main() {
    let map = &Map::from_file("input-day18.txt").unwrap();
    println!("{:?}", count_steps(map));

    let mut steps = 0_usize;
    for partition in map.partition() {
        steps += count_steps(&partition).unwrap();
    }
    println!("{:?}", steps);
}