2022: d16: ex2: add solution

2022/d16/ex1/ex1.py

@@ -0,0 +1,125 @@
+#!/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()

2022/d16/ex1/input

@@ -0,0 +1,61 @@
+Valve WT has flow rate=0; tunnels lead to valves BD, FQ
+Valve UG has flow rate=0; tunnels lead to valves FQ, YB
+Valve FN has flow rate=0; tunnels lead to valves TV, GA
+Valve RU has flow rate=11; tunnels lead to valves YZ, QS, BL, BT, WJ
+Valve RH has flow rate=0; tunnels lead to valves AS, II
+Valve FL has flow rate=0; tunnels lead to valves HR, PQ
+Valve KQ has flow rate=18; tunnels lead to valves FR, BN
+Valve PM has flow rate=25; tunnels lead to valves YZ, FR
+Valve RQ has flow rate=0; tunnels lead to valves FQ, MW
+Valve BL has flow rate=0; tunnels lead to valves RU, IR
+Valve FF has flow rate=0; tunnels lead to valves QS, ED
+Valve KP has flow rate=0; tunnels lead to valves QM, MA
+Valve YB has flow rate=0; tunnels lead to valves UG, HR
+Valve TV has flow rate=17; tunnels lead to valves BD, MT, FN
+Valve HY has flow rate=0; tunnels lead to valves DW, IU
+Valve KF has flow rate=0; tunnels lead to valves AA, HR
+Valve YC has flow rate=0; tunnels lead to valves II, MA
+Valve EE has flow rate=0; tunnels lead to valves AA, CD
+Valve ED has flow rate=9; tunnels lead to valves HG, FF
+Valve SA has flow rate=0; tunnels lead to valves MW, LS
+Valve II has flow rate=20; tunnels lead to valves YC, CY, QP, RH
+Valve BN has flow rate=0; tunnels lead to valves BT, KQ
+Valve MO has flow rate=0; tunnels lead to valves XO, VI
+Valve YZ has flow rate=0; tunnels lead to valves RU, PM
+Valve WJ has flow rate=0; tunnels lead to valves RU, QP
+Valve AW has flow rate=0; tunnels lead to valves HR, DW
+Valve MJ has flow rate=0; tunnels lead to valves BP, AA
+Valve DW has flow rate=4; tunnels lead to valves AU, CB, HY, GL, AW
+Valve QM has flow rate=0; tunnels lead to valves KP, FQ
+Valve LF has flow rate=5; tunnels lead to valves LS, QN, AU, BP, ZY
+Valve QS has flow rate=0; tunnels lead to valves FF, RU
+Valve BT has flow rate=0; tunnels lead to valves BN, RU
+Valve VI has flow rate=22; tunnel leads to valve MO
+Valve LS has flow rate=0; tunnels lead to valves LF, SA
+Valve QD has flow rate=0; tunnels lead to valves HR, ZY
+Valve HG has flow rate=0; tunnels lead to valves AS, ED
+Valve BD has flow rate=0; tunnels lead to valves WT, TV
+Valve CD has flow rate=0; tunnels lead to valves EE, MW
+Valve QP has flow rate=0; tunnels lead to valves II, WJ
+Valve MW has flow rate=7; tunnels lead to valves PQ, SA, CB, CD, RQ
+Valve AU has flow rate=0; tunnels lead to valves DW, LF
+Valve RR has flow rate=0; tunnels lead to valves AS, MA
+Valve GA has flow rate=0; tunnels lead to valves FN, MA
+Valve MT has flow rate=0; tunnels lead to valves CY, TV
+Valve HR has flow rate=14; tunnels lead to valves KF, YB, QD, AW, FL
+Valve AS has flow rate=16; tunnels lead to valves RR, RH, HG, IR
+Valve CY has flow rate=0; tunnels lead to valves MT, II
+Valve AA has flow rate=0; tunnels lead to valves OX, KF, GL, MJ, EE
+Valve IU has flow rate=0; tunnels lead to valves XO, HY
+Valve XO has flow rate=23; tunnels lead to valves IU, MO
+Valve FR has flow rate=0; tunnels lead to valves KQ, PM
+Valve CB has flow rate=0; tunnels lead to valves MW, DW
+Valve ZY has flow rate=0; tunnels lead to valves QD, LF
+Valve BP has flow rate=0; tunnels lead to valves LF, MJ
+Valve QN has flow rate=0; tunnels lead to valves LF, FQ
+Valve IR has flow rate=0; tunnels lead to valves AS, BL
+Valve PQ has flow rate=0; tunnels lead to valves FL, MW
+Valve GL has flow rate=0; tunnels lead to valves AA, DW
+Valve OX has flow rate=0; tunnels lead to valves MA, AA
+Valve MA has flow rate=10; tunnels lead to valves RR, YC, GA, OX, KP
+Valve FQ has flow rate=12; tunnels lead to valves QN, WT, UG, RQ, QM

2022/d16/ex2/ex2.py

@@ -0,0 +1,128 @@
+#!/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]
+        ) -> tuple[int, frozenset[str]]:
+            assert time >= 0  # Sanity check
+            if time == 0:
+                return 0, opened_valves
+
+            pressure = pressure_per_minute(opened_valves)
+
+            # Base-case, don't do anything
+            best = pressure * time, opened_valves
+
+            # Try to open the current valve if not done already
+            if start not in opened_valves:
+                score, valves = helper(start, time - 1, opened_valves | {start})
+                score += pressure
+                best = max(best, (score, valves))
+
+            # Try to go to each neighbour
+            for n, d in dist[start].items():
+                if d >= time:
+                    continue
+                score, valves = helper(n, time - d, opened_valves)
+                score += pressure * d
+                best = max(best, (score, valves))
+
+            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)
+        score, valves = helper(START_ROOM, 26, frozenset(opened_valves))
+        elephant_score, _ = helper(START_ROOM, 26, valves)
+        elephant_score -= pressure_per_minute(valves) * 26  # Don't double count
+        return score + elephant_score
+
+    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()

2022/d16/ex2/input

@@ -0,0 +1,61 @@
+Valve WT has flow rate=0; tunnels lead to valves BD, FQ
+Valve UG has flow rate=0; tunnels lead to valves FQ, YB
+Valve FN has flow rate=0; tunnels lead to valves TV, GA
+Valve RU has flow rate=11; tunnels lead to valves YZ, QS, BL, BT, WJ
+Valve RH has flow rate=0; tunnels lead to valves AS, II
+Valve FL has flow rate=0; tunnels lead to valves HR, PQ
+Valve KQ has flow rate=18; tunnels lead to valves FR, BN
+Valve PM has flow rate=25; tunnels lead to valves YZ, FR
+Valve RQ has flow rate=0; tunnels lead to valves FQ, MW
+Valve BL has flow rate=0; tunnels lead to valves RU, IR
+Valve FF has flow rate=0; tunnels lead to valves QS, ED
+Valve KP has flow rate=0; tunnels lead to valves QM, MA
+Valve YB has flow rate=0; tunnels lead to valves UG, HR
+Valve TV has flow rate=17; tunnels lead to valves BD, MT, FN
+Valve HY has flow rate=0; tunnels lead to valves DW, IU
+Valve KF has flow rate=0; tunnels lead to valves AA, HR
+Valve YC has flow rate=0; tunnels lead to valves II, MA
+Valve EE has flow rate=0; tunnels lead to valves AA, CD
+Valve ED has flow rate=9; tunnels lead to valves HG, FF
+Valve SA has flow rate=0; tunnels lead to valves MW, LS
+Valve II has flow rate=20; tunnels lead to valves YC, CY, QP, RH
+Valve BN has flow rate=0; tunnels lead to valves BT, KQ
+Valve MO has flow rate=0; tunnels lead to valves XO, VI
+Valve YZ has flow rate=0; tunnels lead to valves RU, PM
+Valve WJ has flow rate=0; tunnels lead to valves RU, QP
+Valve AW has flow rate=0; tunnels lead to valves HR, DW
+Valve MJ has flow rate=0; tunnels lead to valves BP, AA
+Valve DW has flow rate=4; tunnels lead to valves AU, CB, HY, GL, AW
+Valve QM has flow rate=0; tunnels lead to valves KP, FQ
+Valve LF has flow rate=5; tunnels lead to valves LS, QN, AU, BP, ZY
+Valve QS has flow rate=0; tunnels lead to valves FF, RU
+Valve BT has flow rate=0; tunnels lead to valves BN, RU
+Valve VI has flow rate=22; tunnel leads to valve MO
+Valve LS has flow rate=0; tunnels lead to valves LF, SA
+Valve QD has flow rate=0; tunnels lead to valves HR, ZY
+Valve HG has flow rate=0; tunnels lead to valves AS, ED
+Valve BD has flow rate=0; tunnels lead to valves WT, TV
+Valve CD has flow rate=0; tunnels lead to valves EE, MW
+Valve QP has flow rate=0; tunnels lead to valves II, WJ
+Valve MW has flow rate=7; tunnels lead to valves PQ, SA, CB, CD, RQ
+Valve AU has flow rate=0; tunnels lead to valves DW, LF
+Valve RR has flow rate=0; tunnels lead to valves AS, MA
+Valve GA has flow rate=0; tunnels lead to valves FN, MA
+Valve MT has flow rate=0; tunnels lead to valves CY, TV
+Valve HR has flow rate=14; tunnels lead to valves KF, YB, QD, AW, FL
+Valve AS has flow rate=16; tunnels lead to valves RR, RH, HG, IR
+Valve CY has flow rate=0; tunnels lead to valves MT, II
+Valve AA has flow rate=0; tunnels lead to valves OX, KF, GL, MJ, EE
+Valve IU has flow rate=0; tunnels lead to valves XO, HY
+Valve XO has flow rate=23; tunnels lead to valves IU, MO
+Valve FR has flow rate=0; tunnels lead to valves KQ, PM
+Valve CB has flow rate=0; tunnels lead to valves MW, DW
+Valve ZY has flow rate=0; tunnels lead to valves QD, LF
+Valve BP has flow rate=0; tunnels lead to valves LF, MJ
+Valve QN has flow rate=0; tunnels lead to valves LF, FQ
+Valve IR has flow rate=0; tunnels lead to valves AS, BL
+Valve PQ has flow rate=0; tunnels lead to valves FL, MW
+Valve GL has flow rate=0; tunnels lead to valves AA, DW
+Valve OX has flow rate=0; tunnels lead to valves MA, AA
+Valve MA has flow rate=10; tunnels lead to valves RR, YC, GA, OX, KP
+Valve FQ has flow rate=12; tunnels lead to valves QN, WT, UG, RQ, QM