Day 16: Proboscidea Volcanium

I'm using a graph library to get the distances between all pairs valves; it just saves me with the tedious busywork of repeated bfs.

(define (part-01 input)
  (define G (pipe-graph input))
  (define paths (find-path G 30))
  (for/fold ([max-flow 0])
            ([path (in-list paths)])
    (max max-flow (flow path))))

(define (pipe-graph input)
  (define pipes (read-valves input))
  ;; valves with flow (plus "AA")
  (define valves (connected-valves pipes))
  ;; all distances
  (define dists (distances pipes))
  (for/hash ([u (in-set valves)])
    (values u
            (cons (first (hash-ref pipes u))
                  (for/list ([v (in-list (neighbours dists valves u))])
                    (list v (hash-ref dists (list u v))))))))

(define (find-path G mins)
  (define valves (list->set (hash-keys G)))
  (define start (list "AA" (- mins 0) 0))
  (define paths null)
  (define (walk u t path opened)
    (cond [(> t mins)]
          [(equal? opened valves)]
          [else
           (define next (rest (hash-ref G u)))
           (for ([vt (in-list next)]
                 #:do [(match-define (list v dt) vt)]
                 #:unless (set-member? opened v))
             (define open-time (+ t dt 1)) ;valve opens
             (define flow-time (- mins open-time))
             (define path+ (cons (list v flow-time (first (hash-ref G v))) path))
             (set! paths (cons path+ paths))
             (walk v open-time path+ (set-add opened v)))]))

  (walk "AA" 0 (list start) (set "AA"))
  paths)

(define (flow path)
  (define start (reverse path))
  (for/fold ([total 0])
            ([u (in-list start)]
             [vopen (in-list (rest start))])
    (match-define (list v flow-time pressure) vopen)
    (+ total (* flow-time pressure))))

(define (neighbours dists valves u)
  (for/list ([edge (in-hash-keys dists)]
             #:when (string=? u (first edge))
             #:unless (string=? u (second edge))
             #:when (set-member? valves (second edge)))
    (second edge)))

(define (connected-valves pipes)
  (set-add
   (for/set ([(k v) (in-hash pipes)] #:when (> (first v) 0)) k)
   "AA"))

(define (distances pipes)
  (define es
    (for*/list ([(from x) (in-hash pipes)]
                [to (in-list (rest x))])
      (list 1 from to)))
  (johnson (weighted-graph/directed es)))

(define (read-valves input)
  (for/hash ([line (in-list input)])
    (match-define (list* _ from (app string->number flow) tos)
      (regexp-match* #px"[A-Z]+|\\d+" line))
    (values from (cons flow tos) )))

I've omitted the code for parsing and pathfinding -- the former is just standard nom parsing and the latter is (a) ugly and (b) a direct translation from Python code on Wikipedia.

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

use ndarray::Array2;
use parser::parse;

mod parser; // Parses the puzzle input
mod seidel; // Computes all-pairs shortest paths, used by parser module to complete the Puzzle struct (https://en.wikipedia.org/wiki/Seidel%27s_algorithm)

#[derive(Debug)]
pub struct Puzzle {
    indices: BTreeMap<String, usize>,
    flow_rate: Vec<usize>,
    adjacency: Array2<usize>,
    path_lengths: Array2<usize>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SearchState {
    unopened: Vec<usize>,
    current_position: usize,
    remaining_time: usize,
    cumulative_flow: usize,
    pressure_released: usize,
}

/// A breadth-first search on the tree implied by the SearchState transitions.
///
/// It prunes branches that cannot beat the current best (ala Branch & Bound).
/// This improves the runtime on my puzzle input from 250ms to 12ms. It solves
/// both the test input and my actual input, but I'm not entirely sure if it is
/// actually sound, or if there are inputs for which it will fail.
fn part1(puzzle: Puzzle) -> usize {
    let mut best = 0;
    let mut q = VecDeque::new();

    let initial_state = SearchState {
        // Filter out valves with flow rate 0, no point in opening them
        unopened: puzzle
            .indices
            .values()
            .cloned()
            .filter(|idx| puzzle.flow_rate[*idx] > 0)
            .collect(),

        current_position: puzzle.indices["AA"],
        remaining_time: 30,
        cumulative_flow: 0,
        pressure_released: 0,
    };

    q.push_back(initial_state);

    while let Some(state) = q.pop_front() {
        if state.unopened.is_empty() || state.remaining_time == 0 {
            // We've ran out of time or valves to open
            best = best.max(state.pressure_released + state.remaining_time * state.cumulative_flow);
            continue;
        }

        for unopened in state.unopened.iter() {
            let mut neighbor = state.clone();

            let movement_time = puzzle.path_lengths[[state.current_position, *unopened]];

            if movement_time >= neighbor.remaining_time {
                // Not enough time left, count out remaining flow
                neighbor.pressure_released += neighbor.remaining_time * neighbor.cumulative_flow;
                neighbor.remaining_time = 0;
            } else {
                // Go to valve, pressure is released while we are en route
                neighbor.remaining_time -= movement_time;
                neighbor.pressure_released += movement_time * neighbor.cumulative_flow;
                neighbor.current_position = *unopened;

                // Open valve, pressure is released while we do it
                neighbor.remaining_time -= 1;
                neighbor.pressure_released += neighbor.cumulative_flow;

                neighbor.unopened.retain(|idx| idx != unopened);
                neighbor.cumulative_flow += puzzle.flow_rate[*unopened];
            }

            if neighbor.pressure_released + neighbor.remaining_time * neighbor.cumulative_flow
                > best
            {
                q.push_back(neighbor);
            }
        }
    }

    best
}

fn part1_input() -> &'static str {
    include_str!("../input.txt")
}

fn main() {
    dbg!(part1(parse(part1_input())));
}