#!/usr/bin/env python import dataclasses import functools import sys from collections import defaultdict @dataclasses.dataclass class Valve: flow: int neighbours: set[str] Graph = dict[str, Valve] DistanceMatrix = dict[str, dict[str, int]] START_ROOM = "AA" def solve(input: list[str]) -> int: def to_graph(input: list[str]) -> Graph: res = {} for line in input: assert line.startswith("Valve ") # Sanity check name = line.split()[1] flow = line.split("=")[1].split(";")[0] neighbours = line.split(";")[1].replace(", ", " ").split()[4:] res[name] = Valve(int(flow), set(neighbours)) return res def useful_valves(g: Graph) -> set[str]: return {k for k, v in g.items() if v.flow > 0} def floyd_warshall(g: Graph) -> DistanceMatrix: points = list(g.keys()) res: DistanceMatrix = defaultdict(dict) for p in points: for n in g[p].neighbours: res[p][n] = 1 for p in points: for i in points: for j in points: if (ip := res[i].get(p)) is None: continue if (pj := res[p].get(j)) is None: continue dist = ip + pj if (ij := res[i].get(j)) is not None: dist = min(dist, ij) res[i][j] = dist return res def prune_distances(dist: DistanceMatrix, of_interest: set[str]) -> DistanceMatrix: # Only keep non-zero valves for our visits pruned = { i: {j: dist for j, dist in line.items() if j in of_interest} for i, line in dist.items() if i in of_interest } # Explicitly add the starting room, in case it was pruned pruned[START_ROOM] = { n: dist for n, dist in dist[START_ROOM].items() if n in of_interest } return pruned def max_flow(g: Graph, dist: DistanceMatrix) -> int: def pressure_per_minute(opened_valves: frozenset[str]) -> int: return sum(g[valve].flow for valve in opened_valves) @functools.cache def helper(start: str, time: int, opened_valves: frozenset[str]) -> int: assert time >= 0 # Sanity check if time == 0: return 0 pressure = pressure_per_minute(opened_valves) # Base-case, don't do anything best = pressure * time # Try to open the current valve if not done already if start not in opened_valves: best = max( best, helper(start, time - 1, opened_valves | {start}) + pressure, ) # Try to go to each neighbour for n, d in dist[start].items(): if d >= time: continue best = max( best, helper(n, time - d, opened_valves) + pressure * d, ) return best opened_valves = set() # If starting room has no flow, consider it open to reduce search space if g[START_ROOM].flow == 0: opened_valves.add(START_ROOM) return helper(START_ROOM, 30, frozenset(opened_valves)) graph = to_graph(input) dist = prune_distances(floyd_warshall(graph), useful_valves(graph)) return max_flow(graph, dist) def main() -> None: input = sys.stdin.read().splitlines() print(solve(input)) if __name__ == "__main__": main()