380 lines
12 KiB
Rust
380 lines
12 KiB
Rust
use std::{
|
|
collections::{HashMap, HashSet, VecDeque},
|
|
hash::Hash,
|
|
};
|
|
|
|
use itertools::Itertools;
|
|
use regex::Regex;
|
|
|
|
use crate::util;
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct Valve {
|
|
pub flow_rate: usize,
|
|
pub tunnels: Vec<usize>,
|
|
}
|
|
|
|
pub fn parse_valves(input: &String) -> Vec<Valve> {
|
|
let indices: HashMap<String, usize> = input
|
|
.lines()
|
|
.map(|l| l.split(' ').nth(1).unwrap())
|
|
.sorted()
|
|
.enumerate()
|
|
.map(|(i, n)| (n.to_owned(), i))
|
|
.collect();
|
|
|
|
let mut valves = HashMap::new();
|
|
|
|
let re =
|
|
Regex::new(r"^Valve (\w+) has flow rate=(\d+); tunnels? leads? to valves? (.*)$").unwrap();
|
|
for line in input.lines() {
|
|
let captures = re
|
|
.captures(line)
|
|
.expect(&format!("No captures for {}", line));
|
|
let name = captures.get(1).unwrap().as_str().to_owned();
|
|
let flow_rate = captures.get(2).unwrap().as_str().parse().unwrap();
|
|
let tunnels = captures
|
|
.get(3)
|
|
.unwrap()
|
|
.as_str()
|
|
.split(", ")
|
|
.map(|s| indices[s])
|
|
.collect();
|
|
let valve = Valve { flow_rate, tunnels };
|
|
valves.insert(name, valve);
|
|
}
|
|
|
|
valves
|
|
.iter()
|
|
.map(|(name, v)| (indices[name], v))
|
|
.sorted_by_key(|(i, _)| *i)
|
|
.map(|(_, v)| v.clone())
|
|
.collect()
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct Path {
|
|
pub dist: usize,
|
|
pub path: Vec<usize>,
|
|
}
|
|
|
|
pub fn calc_paths(valves: &Vec<Valve>) -> Vec<Vec<Path>> {
|
|
let n = valves.len();
|
|
|
|
// result[from][to] = (dist, path)
|
|
let mut result = vec![
|
|
vec![
|
|
Path {
|
|
dist: usize::MAX,
|
|
path: Vec::new()
|
|
};
|
|
n
|
|
];
|
|
n
|
|
];
|
|
|
|
for i in 0..n {
|
|
result[i][i].dist = 0;
|
|
let mut next = VecDeque::new();
|
|
next.push_back(i);
|
|
while let Some(j) = next.pop_front() {
|
|
let d_next = result[i][j].dist + 1;
|
|
for tunnel in &valves[j].tunnels {
|
|
let d_tunnel = &result[i][*tunnel];
|
|
if d_next < d_tunnel.dist {
|
|
let mut path = result[i][j].path.clone();
|
|
path.push(*tunnel);
|
|
result[i][*tunnel] = Path { dist: d_next, path };
|
|
next.push_back(*tunnel);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
result
|
|
}
|
|
|
|
#[derive(Debug, Clone, Hash, PartialEq)]
|
|
pub enum Goal {
|
|
Idle,
|
|
NothingReachable,
|
|
Opening(usize),
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct State {
|
|
pub opened: u64,
|
|
pub locations: Vec<usize>,
|
|
pub goals: Vec<Goal>,
|
|
pub flow_rate: usize,
|
|
pub time_remaining: usize,
|
|
pub pressure_released: usize,
|
|
pub history: Vec<Vec<usize>>,
|
|
|
|
interesting: HashSet<usize>,
|
|
}
|
|
|
|
impl State {
|
|
pub fn new(valves: &Vec<Valve>, initial_time: usize, agents: usize) -> State {
|
|
let interesting = valves
|
|
.iter()
|
|
.enumerate()
|
|
.filter_map(|(i, v)| if v.flow_rate > 0 { Some(i) } else { None })
|
|
.collect();
|
|
State {
|
|
opened: 0,
|
|
locations: vec![0; agents],
|
|
goals: vec![Goal::Idle; agents],
|
|
flow_rate: 0,
|
|
time_remaining: initial_time,
|
|
pressure_released: 0,
|
|
interesting,
|
|
history: vec![vec![0]; agents],
|
|
}
|
|
}
|
|
|
|
fn find_goals(&self, (_, paths): &ExtraArgs) -> Vec<State> {
|
|
let target_dists: HashMap<usize, usize> = self
|
|
.goals
|
|
.iter()
|
|
.enumerate()
|
|
.filter_map(|(i, g)| {
|
|
if let Goal::Opening(t) = g {
|
|
Some((*t, paths[self.locations[i]][*t].dist))
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
.collect();
|
|
|
|
let new_goals: HashMap<usize, Vec<usize>> = self
|
|
.goals
|
|
.iter()
|
|
.positions(|g| matches!(g, Goal::Idle))
|
|
.map(|i| {
|
|
let loc = self.locations[i];
|
|
(
|
|
i,
|
|
self.interesting
|
|
.iter()
|
|
.filter_map(|target| {
|
|
let agent_dist = paths[loc][*target].dist;
|
|
if agent_dist + 2 >= self.time_remaining {
|
|
// Can't reach + open the valve as well as
|
|
// letting anything flow through
|
|
return None;
|
|
}
|
|
|
|
if let Some(dist) = target_dists.get(target) {
|
|
if *dist > agent_dist {
|
|
// Someone else is targeting this, but they're slower
|
|
Some(*target)
|
|
} else {
|
|
// Someone else is targeting this, and they're faster
|
|
None
|
|
}
|
|
} else {
|
|
// Noone else is targeting this
|
|
Some(*target)
|
|
}
|
|
})
|
|
.collect_vec(),
|
|
)
|
|
})
|
|
.collect();
|
|
|
|
if new_goals.len() == 1 {
|
|
let (i, targets) = new_goals.iter().next().unwrap();
|
|
if targets.is_empty() {
|
|
let mut state = self.clone();
|
|
state.goals[*i] = Goal::NothingReachable;
|
|
return vec![state];
|
|
}
|
|
|
|
targets
|
|
.iter()
|
|
.map(|t| {
|
|
let mut state = self.clone();
|
|
state.goals[*i] = Goal::Opening(*t);
|
|
state
|
|
})
|
|
.collect()
|
|
} else if new_goals.len() == 2 {
|
|
let mut it = new_goals.iter();
|
|
let (i1, targets1) = it.next().unwrap();
|
|
let (i2, targets2) = it.next().unwrap();
|
|
if targets1.len() == 0 {
|
|
if targets2.len() == 0 {
|
|
let mut state = self.clone();
|
|
state.goals[*i1] = Goal::NothingReachable;
|
|
state.goals[*i2] = Goal::NothingReachable;
|
|
return vec![state];
|
|
}
|
|
|
|
return targets2
|
|
.iter()
|
|
.map(|t| {
|
|
let mut state = self.clone();
|
|
state.goals[*i1] = Goal::NothingReachable;
|
|
state.goals[*i2] = Goal::Opening(*t);
|
|
state
|
|
})
|
|
.collect();
|
|
} else if targets2.len() == 0 {
|
|
return targets1
|
|
.iter()
|
|
.map(|t| {
|
|
let mut state = self.clone();
|
|
state.goals[*i1] = Goal::Opening(*t);
|
|
state.goals[*i2] = Goal::NothingReachable;
|
|
state
|
|
})
|
|
.collect();
|
|
} else if targets1.len() == 1 && targets1 == targets2 {
|
|
// Both agents can only reach one and the same goal
|
|
let mut state = self.clone();
|
|
let target = targets1[0];
|
|
if paths[self.locations[*i1]][target].dist < paths[self.locations[*i2]][target].dist
|
|
{
|
|
state.goals[*i1] = Goal::Opening(target);
|
|
state.goals[*i2] = Goal::NothingReachable;
|
|
} else {
|
|
state.goals[*i1] = Goal::NothingReachable;
|
|
state.goals[*i2] = Goal::Opening(target);
|
|
}
|
|
return vec![state];
|
|
}
|
|
|
|
let mut result = Vec::new();
|
|
|
|
for target1 in targets1 {
|
|
for target2 in targets2 {
|
|
if target1 == target2 {
|
|
continue;
|
|
}
|
|
|
|
let mut state = self.clone();
|
|
state.goals[*i1] = Goal::Opening(*target1);
|
|
state.goals[*i2] = Goal::Opening(*target2);
|
|
result.push(state);
|
|
}
|
|
}
|
|
|
|
result
|
|
} else {
|
|
unimplemented!("Can't find goals for >2 agents at the same time");
|
|
}
|
|
}
|
|
}
|
|
|
|
type ExtraArgs<'a> = (&'a Vec<Valve>, &'a Vec<Vec<Path>>);
|
|
impl util::BnBState<ExtraArgs<'_>> for State {
|
|
fn possible_actions(&self, (valves, paths): &ExtraArgs) -> Vec<State> {
|
|
assert!(!self.finished());
|
|
|
|
let goalless = self
|
|
.goals
|
|
.iter()
|
|
.enumerate()
|
|
.filter(|(_, g)| matches!(g, Goal::Idle))
|
|
.collect_vec();
|
|
if goalless.len() > 0 {
|
|
return self.find_goals(&(valves, paths));
|
|
}
|
|
|
|
let min_goal_dist = self
|
|
.goals
|
|
.iter()
|
|
.enumerate()
|
|
.filter_map(|(i, g)| {
|
|
if let Goal::Opening(target) = g {
|
|
Some(paths[self.locations[i]][*target].dist)
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
.min()
|
|
.unwrap();
|
|
|
|
let time_spent = min_goal_dist + 1;
|
|
let mut completed = self.clone();
|
|
completed.pressure_released += completed.flow_rate * time_spent;
|
|
completed.time_remaining -= time_spent;
|
|
for (i, goal) in self.goals.iter().enumerate() {
|
|
if let Goal::Opening(goal) = goal {
|
|
let path = &paths[self.locations[i]][*goal];
|
|
if path.dist == min_goal_dist {
|
|
completed.locations[i] = *goal;
|
|
completed.opened |= 1 << goal;
|
|
completed.goals[i] = Goal::Idle;
|
|
if completed.interesting.remove(goal) {
|
|
// Only increase flow rate if no other agent already
|
|
// opened this valve
|
|
completed.flow_rate += valves[*goal].flow_rate;
|
|
}
|
|
completed.history[i].push(*goal);
|
|
} else {
|
|
completed.locations[i] = path.path[time_spent - 1];
|
|
}
|
|
}
|
|
}
|
|
return vec![completed];
|
|
}
|
|
|
|
fn finished(&self) -> bool {
|
|
self.time_remaining == 0
|
|
|| self
|
|
.goals
|
|
.iter()
|
|
.all(|g| matches!(g, Goal::NothingReachable)) //self.interesting.is_empty()
|
|
}
|
|
|
|
fn lower_bound(&self, _: &ExtraArgs) -> usize {
|
|
self.pressure_released + self.flow_rate * self.time_remaining
|
|
}
|
|
|
|
fn upper_bound(&self, (valves, paths): &ExtraArgs) -> usize {
|
|
let mut additional_flow = 0;
|
|
for i in &self.interesting {
|
|
let i = i.clone();
|
|
if (self.opened & (1 << i)) != 0 {
|
|
continue;
|
|
}
|
|
let dist = self
|
|
.locations
|
|
.iter()
|
|
.map(|loc| paths[*loc][i].dist)
|
|
.min()
|
|
.unwrap();
|
|
if dist < self.time_remaining {
|
|
let flow = valves[i].flow_rate;
|
|
additional_flow += flow * (self.time_remaining - dist - 1);
|
|
}
|
|
}
|
|
|
|
self.lower_bound(&(valves, paths)) + additional_flow
|
|
}
|
|
}
|
|
|
|
impl Hash for State {
|
|
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
|
|
self.opened.hash(state);
|
|
self.locations.hash(state);
|
|
self.goals.hash(state);
|
|
self.time_remaining.hash(state);
|
|
self.pressure_released.hash(state);
|
|
}
|
|
}
|
|
|
|
impl PartialEq for State {
|
|
fn eq(&self, other: &Self) -> bool {
|
|
self.opened == other.opened
|
|
&& self.locations == other.locations
|
|
&& self.goals == other.goals
|
|
&& self.time_remaining == other.time_remaining
|
|
&& self.pressure_released == other.pressure_released
|
|
}
|
|
}
|
|
|
|
impl Eq for State {}
|