advent_solution_day16_v2.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 positions: [u8; 2],
}

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(),
        positions: [start_idx, start_idx],
    };
    let mut states: HashMap<State, i32> = HashMap::new();
    states.insert(start_state, 0);
    for minute in (1..26).rev() {
        println!("{}",minute);
        for actor_idx in 0..2{
            let mut new_states: HashMap<State, i32> = HashMap::new();
            for (old_state, old_flow) in states.iter() {
                let old_pos = old_state.positions[actor_idx] as usize;
                let mut bitmap_flipped = old_state.bitmap;
                bitmap_flipped.set(old_pos, true);
                for new_pos in network[old_pos].iter() {
                    let mut new_posses = old_state.positions;
                    new_posses[actor_idx] = *new_pos;
                    let new_state = State {
                        bitmap: old_state.bitmap,
                        positions: new_posses,
                    };
                    let entry = new_states.entry(new_state).or_insert(*old_flow);
                    *entry = max(*entry, *old_flow);
                }
                if !old_state.bitmap.get(old_pos) && flows[old_pos] > 0{
                    let mut bitmap_flipped = old_state.bitmap;
                    bitmap_flipped.set(old_pos, true);
                    let new_flow = *old_flow + flows[old_pos] * minute;
                    let new_state = State {
                        bitmap: bitmap_flipped,
                        positions: old_state.positions,
                    };
                    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);
}