advent_solution_day16_v1.py

use bitmaps::Bitmap;
use std::cmp::max;
use std::collections::HashMap;
use std::hash::Hash;
use std::io;

fn parsei(s: &str) -> i32 {
    s.parse().unwrap()
}
const NUM_LETTERS: usize = 26;
// const MAX_NUM_VALVES: usize = NUM_LETTERS * NUM_LETTERS;
fn parse_valve(s: &str) -> usize {
    let bytes = s.as_bytes();
    (((bytes[0] - b'A') as usize) * NUM_LETTERS) + (bytes[1] - b'A') as usize
}
type FlowMask = Bitmap<50>;

#[derive(PartialEq, Eq, Debug, Hash, Copy, Clone)]
struct State {
    pub bitmap: FlowMask,
    pub position: u8,
}

fn main() {
    let stdin = io::stdin();

    let mut valves: Vec<usize> = Vec::new();
    let mut network_f: Vec<Vec<usize>> = Vec::new();
    let mut flows: Vec<i32> = Vec::new();
    loop {
        let mut user_input = String::new();
        let size = stdin.read_line(&mut user_input).unwrap();
        if user_input.len() == 0 {
            break;
        }
        let items: Vec<&str> = user_input.trim().split_ascii_whitespace().collect();
        let valve_name = parse_valve(items[1]);
        let flow_rate = parsei(
            items[4]
                .strip_prefix("rate=")
                .unwrap()
                .strip_suffix(";")
                .unwrap(),
        );
        let valve_dests = &items[9..];
        let parsed_valve_dests: Vec<usize> = valve_dests
            .iter()
            .map(|v| parse_valve(v.strip_suffix(",").unwrap_or(v)))
            .collect();
        network_f.push(parsed_valve_dests);
        flows.push(flow_rate);
        valves.push(valve_name);
    }
    let mut rev_valve_map = [0 as u8; u16::MAX as usize];
    for (i, item) in valves.iter().enumerate() {
        rev_valve_map[*item] = i as u8;
    }
    let start_idx = rev_valve_map[parse_valve("AA")];
    let network: Vec<Vec<u8>> = network_f
        .iter()
        .map(|n| n.iter().map(|v| rev_valve_map[*v]).collect())
        .collect();
    let start_state = State {
        bitmap: FlowMask::new(),
        position: start_idx,
    };
    let mut states: HashMap<State, i32> = HashMap::new();
    states.insert(start_state, 0);
    for minute in (1..30).rev() {
        let mut new_states: HashMap<State, i32> = HashMap::new();
        for (old_state, old_flow) in states.iter() {
            let mut bitmap_flipped = old_state.bitmap;
            bitmap_flipped.set(old_state.position as usize, true);
            for new_pos in network[old_state.position as usize].iter() {
                let new_state = State {
                    bitmap: old_state.bitmap,
                    position: *new_pos,
                };
                let entry = new_states.entry(new_state).or_insert(*old_flow);
                *entry = max(*entry, *old_flow);
            }
            if !old_state.bitmap.get(old_state.position as usize) && flows[old_state.position as usize] > 0{
                let mut bitmap_flipped = old_state.bitmap;
                bitmap_flipped.set(old_state.position as usize, true);
                let new_flow = *old_flow + flows[old_state.position as usize] * minute;
                let new_state = State {
                    bitmap: bitmap_flipped,
                    position: old_state.position,
                };
                let entry = new_states.entry(new_state).or_insert(new_flow);
                *entry = max(*entry, new_flow);
            }
        }
        states = new_states;
    }
    let max_flow: i32 = states.values().copied().max().unwrap();
    
    println!("{}",max_flow);
}