Files
advent-of-code-2022-rust/src/day16.rs

303 lines
12 KiB
Rust

use std::collections::HashMap;
use crate::day16::ValveStep::{MoveTo, OpenValve};
use crate::day_solver::DaySolver;
use super::util;
pub struct Day16 {
valves: Vec<Valve>,
initial_valve_idx: usize,
valve_names: Vec<String>
}
impl Day16 {
pub fn create() -> Self {
// let lines = util::read_file("input/day16_example.txt");
let lines = util::read_file("input/day16.txt");
// We first map all the valve names to indexes for performance reasons:
let valve_names = lines.iter()
.map(|s| s.split_whitespace().skip(1).next().unwrap().to_string())
.collect::<Vec<String>>();
// Put the input into the day struct
return Day16 {
valves: lines.iter().enumerate().map(|(i, s)| {
let mut s_split = s.split(";");
let flow_rate = s_split.next().unwrap().split("rate=").skip(1).next().unwrap().parse::<u32>().unwrap();
let connections = (&s_split.next().unwrap()["tunnels lead to valves ".len()..].trim()).split(", ")
.map(|n| valve_names.iter().position(|nn| n.eq(nn)).unwrap())
.collect();
Valve {
id: i, flow_rate, connections
}
}).collect(),
initial_valve_idx: valve_names.iter().position(|n| n == "AA").unwrap(),
valve_names
}
}
fn set_open(valve_idx: &usize, cur_open_valves_mask: &u64) -> u64 {
cur_open_valves_mask.clone() | (1u64 << valve_idx)
}
fn is_open(valve_idx: &usize, open_valves_mask: &u64) -> bool {
open_valves_mask & (1u64 << valve_idx) != 0
}
fn is_closed(valve_idx: &usize, open_valves_mask: &u64) -> bool {
!Self::is_open(valve_idx, open_valves_mask)
}
fn calc_flow_rate(&self, open_valves_mask: &u64) -> u32 {
let mut res = 0;
for i in 0..self.valves.len() {
if Self::is_open(&i, open_valves_mask) {
res += self.valves[i].flow_rate
}
}
res
}
fn should_valve_be_opened(&self, valve_idx: &usize, open_valves_mask: &u64) -> bool {
if Self::is_open(valve_idx, open_valves_mask) { false }
else if let Some(valve) = self.valves.get(*valve_idx) {
valve.flow_rate > 0
} else { false }
}
fn find_best_option_pt2(&self, valve_state: ValveStatePt2, cur_flow_rate: u32, cache: &mut HashMap<ValveStatePt2, BestOptionPt2>) -> BestOptionPt2 {
if valve_state.minutes_left == 0 {
return BestOptionPt2 {
// my_step: ValveStep::None,
// elephant_step: ValveStep::None,
total_flow: 0,
// cur_flow: cur_flow_rate
}
}
if let Some(res) = cache.get(&valve_state) {
return res.to_owned();
}
let mut my_options = Vec::new();
let mut elephant_options = Vec::new();
if self.should_valve_be_opened(&valve_state.my_position, &valve_state.open_valves_mask) {
// One option is to open the valve at the current position
my_options.push(OpenValve(valve_state.my_position));
}
if valve_state.my_position != valve_state.elephant_position && self.should_valve_be_opened(&valve_state.elephant_position, &valve_state.open_valves_mask) {
elephant_options.push(OpenValve(valve_state.elephant_position));
}
// We try and move to an adjacent valve, see if we can be more useful there
let my_valve = &self.valves[valve_state.my_position];
for connection in my_valve.connections.iter() {
my_options.push(MoveTo(connection.to_owned()))
}
let elephant_valve = &self.valves[valve_state.elephant_position];
for connection in elephant_valve.connections.iter() {
elephant_options.push(MoveTo(connection.to_owned()))
}
// Now we try all combinations of my and elephants options:
let mut best: u32 = 0;
let mut my_best_step = None;
let mut elephant_best_step = None;
let next_minutes_left = valve_state.minutes_left - 1;
for my_option in my_options {
let mut my_next_position = valve_state.my_position;
let mut next_valve_mask_for_me = valve_state.open_valves_mask;
let mut next_flow_rate_for_me = cur_flow_rate;
match my_option {
MoveTo(idx) => my_next_position = idx,
OpenValve(idx) => {
next_valve_mask_for_me = Self::set_open(&idx, &next_valve_mask_for_me);
next_flow_rate_for_me += self.valves[idx].flow_rate
}
_ => {}
}
for elephant_option in &elephant_options {
let mut next_valve_mask = next_valve_mask_for_me;
let mut next_flow_rate = next_flow_rate_for_me;
let mut elephant_next_position = valve_state.elephant_position;
match elephant_option {
MoveTo(idx) => elephant_next_position = *idx,
OpenValve(idx) => {
next_valve_mask = Self::set_open(&idx, &next_valve_mask);
next_flow_rate += self.valves[*idx].flow_rate
}
_ => {}
}
let res = self.find_best_option_pt2(ValveStatePt2::new(my_next_position, elephant_next_position, next_valve_mask, next_minutes_left), next_flow_rate, cache);
if best == 0 || res.total_flow > best {
best = res.total_flow;
my_best_step = Some(my_option);
elephant_best_step = Some(elephant_option.to_owned());
}
}
}
let res = BestOptionPt2 {
// my_step: my_best_step.unwrap(),
// elephant_step: elephant_best_step.unwrap(),
// cur_flow: cur_flow_rate,
total_flow: best + cur_flow_rate
};
cache.insert(valve_state, res.to_owned());
res
}
fn find_best_option(&self, valve_state: ValveState, cur_flow_rate: u32, cache: &mut HashMap<ValveState, BestOptionResult>) -> Option<BestOptionResult> {
if valve_state.minutes_left == 0 { return None }
if let Some(res) = cache.get(&valve_state) {
return Some(res.clone());
}
let cur_valve = &self.valves.get(valve_state.my_position).unwrap();
let mut best_step_option: Option<ValveStep> = None;
let mut best_step_result_option: Option<BestOptionResult> = None;
if self.should_valve_be_opened(&valve_state.my_position, &valve_state.open_valves_mask) {
// One option is to open the valve at the current position
let new_open_valves_mask = Self::set_open(&valve_state.my_position, &valve_state.open_valves_mask);
best_step_result_option = self.find_best_option(ValveState::new(valve_state.my_position, new_open_valves_mask, valve_state.minutes_left - 1), cur_flow_rate + cur_valve.flow_rate, cache);
best_step_option = Some(ValveStep::OpenValve(valve_state.my_position));
}
// We try and move to an adjacent valve, see if we can be more useful there
for connection in cur_valve.connections.iter().rev() {
let next_step_result = self.find_best_option(ValveState::new(connection.to_owned(), valve_state.open_valves_mask, valve_state.minutes_left - 1), cur_flow_rate, cache);
if best_step_result_option.is_none() || next_step_result.as_ref().unwrap().total_flow > best_step_result_option.as_ref().unwrap().total_flow {
best_step_result_option = next_step_result;
best_step_option = Some(ValveStep::MoveTo(connection.to_owned()));
}
}
let Some(best_step) = best_step_option else { panic!("Couldn't find a good step?") };
let mut flow_per_step = Vec::with_capacity(valve_state.minutes_left as usize);
flow_per_step.push(cur_flow_rate);
let mut steps = Vec::with_capacity(valve_state.minutes_left as usize);
steps.push(best_step);
let mut best_next_flow = 0;
if let Some(best_step_result) = best_step_result_option {
flow_per_step.extend(best_step_result.flow_per_step.iter());
steps.extend(best_step_result.steps.iter());
best_next_flow = best_step_result.total_flow;
}
let step_result = BestOptionResult {
total_flow: best_next_flow + cur_flow_rate,
flow_per_step,
steps,
};
cache.insert(valve_state, step_result.clone());
Some(step_result)
}
}
impl DaySolver for Day16 {
fn solve_part1(&mut self) -> String {
let mut cache= HashMap::new();
// let example_flows = [0, 0, 0, 20, 20, 20, 33, 33, 33, 33, 54, 54, 54, 54, 54, 54, 54, 54, 76, 76, 76, 76, 79, 79, 79, 81, 81, 81, 81, 81, 81];
// let mut prev_best: u32 = 0;
// let mut example_running_total = 0;
// for i in 1..31u32 {
// let Some(best) = self.find_best_option(ValveState::new(self.initial_valve_idx, 0, i), 0, &mut cache) else { panic!("Not best!")};
// let cur_example_flow = example_flows[i as usize];
// example_running_total += cur_example_flow;
// println!("Best in {} minutes is {}, so I guess the flow is {}, while example flow is {} (example running total {})", i, best.total_flow, best.total_flow - prev_best, cur_example_flow, example_running_total);
// prev_best = best.total_flow;
// }
let best = self.find_best_option(ValveState::new(self.initial_valve_idx, 0, 30), 0, &mut cache).unwrap();
// for i in 0..best.steps.len() {
// println!("Step {}", i + 1);
// match best.steps[i] {
// ValveStep::OpenValve(v) => println!("Opening valve {}", self.valve_names[v]),
// ValveStep::MoveTo(v ) => println!("Moving to valve {}", self.valve_names[v])
// }
// println!("The flow is now {}\n", best.flow_per_step[i]);
// }
best.total_flow.to_string()
}
fn solve_part2(&mut self) -> String {
let mut cache= HashMap::new();
let best = self.find_best_option_pt2(ValveStatePt2::new(self.initial_valve_idx, self.initial_valve_idx, 0, 26), 0, &mut cache);
return best.total_flow.to_string();
}
}
#[derive(Debug, Clone)]
struct Valve {
id: usize,
flow_rate: u32,
connections: Vec<usize>
}
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
struct ValveState {
my_position: usize,
open_valves_mask: u64,
minutes_left: u32
}
impl ValveState {
fn new(cur_valve_idx: usize, open_valves_mask: u64, minutes_left: u32) -> Self{
ValveState { my_position: cur_valve_idx, open_valves_mask, minutes_left }
}
}
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
struct ValveStatePt2 {
my_position: usize,
elephant_position: usize,
open_valves_mask: u64,
minutes_left: u32
}
#[derive(Debug, Clone, Copy)]
struct BestOptionPt2 {
// my_step: ValveStep,
// elephant_step: ValveStep,
total_flow: u32,
// cur_flow: u32
}
impl ValveStatePt2 {
fn new(my_position: usize, elephant_position: usize, open_valves_mask: u64, minutes_left: u32) -> Self{
ValveStatePt2 { my_position, elephant_position, open_valves_mask, minutes_left }
}
}
#[derive(Debug, Clone)]
struct BestOptionResult {
steps: Vec<ValveStep>,
flow_per_step: Vec<u32>,
total_flow: u32
}
#[derive(Debug, Clone, Copy, PartialEq)]
enum ValveStep {
OpenValve(usize),
MoveTo(usize),
None
}