Everybody Codes Quest 15 Part 3

quest15.rs

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

const WALL: u8 = 128;
const EMPTY: u8 = 64;

const POS_MASK: u64 = 0x00000000ffffffff;
const HERB_MASK: u64 = 0xffffffff00000000;

struct Grid {
    width: u64,
    height: u64,
    start: u64,
    herb_mask: u64,
    data: Vec<u8>,
}

struct Node {
    state: u64,
    cost: u32
}

fn find_start(first_row: &[u8]) -> usize {
    for (i, c) in first_row.iter().enumerate() {
        if *c == EMPTY {
            return i;
        }
    }
    panic!("Failed to find entrance");
}

fn parse_input() -> Grid {
    let mut data = std::fs::read("input.txt").unwrap();
    let width = data.iter().position(|&c| c == b'\n').unwrap();
    data.retain(|&c| c != b'\n');
    let height = data.len() / width;
    assert!(width * height == data.len());
    let mut herb_types: HashSet<u8> = HashSet::new();
    for c in data.iter_mut() {
        *c = match *c {
            b'.' => EMPTY,
            b'#' | b'~' => WALL,
            b'A'..=b'Z' => {
                let h = *c - b'A';
                herb_types.insert(h);
                h
            },
            _ => panic!("Invalid char in input: {}", *c as char)
        };
    }
    let start = find_start(&data[0..width]);
    let mut herb_mask = 0;
    for h in herb_types {
        herb_mask |= 1 << (h + 32);
    }
    Grid {
        width: width.try_into().unwrap(),
        height: height.try_into().unwrap(),
        start: start.try_into().unwrap(),
        herb_mask,
        data,
    }
}

fn maybe_add_neighbor(grid: &Grid, pos: u64, cost: u32, herbs_collected: u64, visited: &mut HashSet<u64>, queue: &mut VecDeque<Node>) {
    let value = grid.data[pos as usize];
    if value == WALL {
        return;
    }
    let herb = if value == EMPTY {
        0
    } else {
        1 << (value + 32)
    };
    let neighbor = Node {
        state: herb | herbs_collected | pos,
        cost
    };
    if visited.insert(neighbor.state) {
        queue.push_back(neighbor);
    }
}

fn bfs(grid: &Grid) -> u32 {
    let end = grid.herb_mask | grid.start;

    let mut queue = VecDeque::new();
    queue.push_back(Node{state: grid.start, cost: 0});
    let mut visited = HashSet::new();
    visited.insert(grid.start);

    while let Some(node) = queue.pop_front() {
        if node.state == end {
            return node.cost;
        }
        let pos = node.state & POS_MASK;
        let herbs_collected = node.state & HERB_MASK;
        let x = pos % grid.width;
        let y = pos / grid.width;
        if x > 0 {
            maybe_add_neighbor(&grid, pos - 1, node.cost + 1, herbs_collected, &mut visited, &mut queue);
        }
        if x < grid.width - 1 {
            maybe_add_neighbor(&grid, pos + 1, node.cost + 1, herbs_collected, &mut visited, &mut queue);
        }
        if y > 0 {
            maybe_add_neighbor(&grid, pos - grid.width, node.cost + 1, herbs_collected, &mut visited, &mut queue);
        }
        if y < grid.height - 1 {
            maybe_add_neighbor(&grid, pos + grid.width, node.cost + 1, herbs_collected, &mut visited, &mut queue);
        }
    }
    panic!("Failed to find end: {}", end);
}

fn main() {
    let grid = parse_input();
    let answer = bfs(&grid);
    println!("Answer {answer}");
}