aoc-2022-day16.rb

def run input
  $edges = Hash.new { |h,k| h[k] = Hash.new 999 }
  $nodes = {}
  input.split("\n").each do |line|
    name, flow, *children = line.scan /[A-Z]{2}|\d+/
    $nodes[name] = flow.to_i
    children.each { |child| $edges[name][child] = 1 }
  end
  $nodes.keys.each { |node| $edges[node][node] = 0 }
  # build all-pairs shortest distance between all nodes (floyd-warshall)
  $nodes.keys.product($nodes.keys, $nodes.keys).each { |k, i, j|
    $edges[i][j] = [$edges[i][j], $edges[i][k] + $edges[k][j]].min
  }
  $nodes.reject! { |_, flow| flow == 0 } # ignore tunnels with no valves
  # use highest value valves to calculate best case estimates
  $nodes = $nodes.entries.sort_by(&:last).reverse.to_h

  # prune edges and add 1 to costs to account for turning valve time
  $edges.each { |node, h|
    h.select! { |k, v| k != node && ($nodes.include?(k) || k == 'AA') }
    h.transform_values! &:succ
  }

  puts "one agent with 30m can release #{search [['AA', 30]]} pressure"
  puts "two agents with 26m can release #{search [['AA', 26]] * 2} pressure"
end

$best = 0
$cache = {}
# DFS search args: agent positions + remaining time, opened valves, 
# and released pressure from those opened valves over the whole time period.
# returns the pressure that can be released from that state onwards
def search(pos, opened = [], start_pressure = 0)
  $cache[[pos, opened]] ||= (
    remaining = $nodes.keys - opened

    # try to estimate our best possible pressure from visiting all remaining
    # valves, using the current minimum edge length emanating from each 
    times = pos.map &:last
    valid_nodes = pos.map(&:first) + remaining
    estimate = remaining.sum { |child|
      idx = times.index times.max
      times[idx] -= $edges[child].values_at(*valid_nodes).min
      $nodes[child] * [times[idx], 0].max
    }
    # and terminate this subtree if we cannot outperform our best
    return -9999 if start_pressure + estimate <= $best

    # try moving either agent
    moves = pos.permutation.flat_map { |_pos|
        node, time = _pos.first
        remaining.map { |dest|
          _time = time - $edges[node][dest]
          _pos = _pos.clone
          _pos[0] = [dest, _time]
          [$nodes[dest] * _time, dest, _pos, _time]
        }
      }
      .select { |move| move.last > 0 }
      .sort_by! {|move| -move.last } # favoring moves that save the most time
      .map { |new_pressure, dest, _pos|
        new_pressure + search(
          _pos.sort, (opened + [dest]).sort, start_pressure + new_pressure)
      }
    
    (moves.max || 0).tap { |best_move|
      $best = [$best, start_pressure + best_move].max
    }
  )
end